Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Stackless [WIP] #23299

Draft
wants to merge 204 commits into
base: main
Choose a base branch
from
Draft

Stackless [WIP] #23299

wants to merge 204 commits into from

Commits on Jul 8, 2024

  1. controlled burn

    @dougalm
    dougalm committed Jul 8, 2024
    Configuration menu
    Copy the full SHA
    cf94b57 View commit details
    Browse the repository at this point in the history

Commits on Jul 9, 2024

  1. CoreTest.test_jit passing

    @dougalm
    dougalm committed Jul 9, 2024
    Configuration menu
    Copy the full SHA
    1ae1fe9 View commit details
    Browse the repository at this point in the history

  2. CoreTest.test_jvp passing

    @dougalm
    dougalm committed Jul 9, 2024
    Configuration menu
    Copy the full SHA
    2a930e0 View commit details
    Browse the repository at this point in the history

  3. Most of the rest of core tests passing

    @dougalm
    dougalm committed Jul 9, 2024
    Configuration menu
    Copy the full SHA
    06d48fe View commit details
    Browse the repository at this point in the history

  4. more

    @dougalm
    dougalm committed Jul 9, 2024
    Configuration menu
    Copy the full SHA
    ec60fb0 View commit details
    Browse the repository at this point in the history

Commits on Jul 10, 2024

  1. vmap

    @dougalm
    dougalm committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    5825377 View commit details
    Browse the repository at this point in the history

  2. More vmap

    @dougalm
    dougalm committed Jul 10, 2024
    Configuration menu
    Copy the full SHA
    c806225 View commit details
    Browse the repository at this point in the history

Commits on Jul 15, 2024

  1. Some custom vjp/jvp

    @dougalm
    dougalm committed Jul 15, 2024
    Configuration menu
    Copy the full SHA
    1f960d8 View commit details
    Browse the repository at this point in the history

Commits on Jul 16, 2024

  1. Batching custom vjp

    @dougalm
    dougalm committed Jul 16, 2024
    Configuration menu
    Copy the full SHA
    f848128 View commit details
    Browse the repository at this point in the history

  2. More custom vjp tests

    @dougalm
    dougalm committed Jul 16, 2024
    Configuration menu
    Copy the full SHA
    6099f29 View commit details
    Browse the repository at this point in the history

  3. more custom jvp

    @dougalm
    dougalm committed Jul 16, 2024
    Configuration menu
    Copy the full SHA
    38fae8f View commit details
    Browse the repository at this point in the history

Commits on Jul 17, 2024

  1. some control flow tests

    @dougalm
    dougalm committed Jul 17, 2024
    Configuration menu
    Copy the full SHA
    fdb6975 View commit details
    Browse the repository at this point in the history

  2. Remove custom bind

    @dougalm
    dougalm committed Jul 17, 2024
    Configuration menu
    Copy the full SHA
    bcedd70 View commit details
    Browse the repository at this point in the history

  3. more control flow tests

    @dougalm
    dougalm committed Jul 17, 2024
    Configuration menu
    Copy the full SHA
    8f0d867 View commit details
    Browse the repository at this point in the history

Commits on Jul 19, 2024

  1. WIP - mattjj to fix :)

    @dougalm
    dougalm committed Jul 19, 2024
    Configuration menu
    Copy the full SHA
    bfc9f9e View commit details
    Browse the repository at this point in the history

Commits on Jul 20, 2024

  1. fix all_gather on constants

    @mattjj
    mattjj committed Jul 20, 2024
    Configuration menu
    Copy the full SHA
    e86ad4a View commit details
    Browse the repository at this point in the history

Commits on Jul 21, 2024

  1. fix up some confusing logic in all-reduce primitives

    @mattjj
    mattjj committed Jul 21, 2024
    Configuration menu
    Copy the full SHA
    a3ddef6 View commit details
    Browse the repository at this point in the history

  2. update all_to_all batcher to take axis_data

    @mattjj
    mattjj committed Jul 21, 2024
    Configuration menu
    Copy the full SHA
    a16ff7b View commit details
    Browse the repository at this point in the history

  3. one failing test left!

    @mattjj
    mattjj committed Jul 21, 2024
    Configuration menu
    Copy the full SHA
    8444027 View commit details
    Browse the repository at this point in the history

Commits on Jul 24, 2024

  1. progress on control_flow_test.py

    @mattjj
    mattjj committed Jul 24, 2024
    Configuration menu
    Copy the full SHA
    57a4bac View commit details
    Browse the repository at this point in the history

  2. attrs

    @dougalm
    dougalm committed Jul 24, 2024
    Configuration menu
    Copy the full SHA
    ecbe339 View commit details
    Browse the repository at this point in the history

  3. checkify

    @dougalm
    dougalm committed Jul 24, 2024
    Configuration menu
    Copy the full SHA
    87a7c2b View commit details
    Browse the repository at this point in the history

  4. more custom jvp

    @dougalm
    dougalm committed Jul 24, 2024
    Configuration menu
    Copy the full SHA
    a687f16 View commit details
    Browse the repository at this point in the history

Commits on Jul 25, 2024

  1. Merge remote-tracking branch 'origin/main' into stackless

    @mattjj
    mattjj committed Jul 25, 2024
    Configuration menu
    Copy the full SHA
    2483cda View commit details
    Browse the repository at this point in the history

  2. fix

    @mattjj
    mattjj committed Jul 25, 2024
    Configuration menu
    Copy the full SHA
    f3aaa14 View commit details
    Browse the repository at this point in the history

  3. more

    @dougalm
    dougalm committed Jul 25, 2024
    Configuration menu
    Copy the full SHA
    21627e9 View commit details
    Browse the repository at this point in the history

Commits on Jul 26, 2024

  1. short-circuit psum path and some pmap tests passing

    @dougalm
    dougalm committed Jul 26, 2024
    Configuration menu
    Copy the full SHA
    b45ca67 View commit details
    Browse the repository at this point in the history

  2. more pmap

    @dougalm
    dougalm committed Jul 26, 2024
    Configuration menu
    Copy the full SHA
    753feb8 View commit details
    Browse the repository at this point in the history

Commits on Jul 27, 2024

  1. Merge remote-tracking branch 'origin/main' into stackless

    @mattjj
    mattjj committed Jul 27, 2024
    Configuration menu
    Copy the full SHA
    22f6f63 View commit details
    Browse the repository at this point in the history

  2. fix bug introduced in merge

    @mattjj
    mattjj committed Jul 27, 2024
    Configuration menu
    Copy the full SHA
    681f70e View commit details
    Browse the repository at this point in the history

  3. fix custom_solve batching rule with axis_data

    @mattjj
    mattjj committed Jul 27, 2024
    Configuration menu
    Copy the full SHA
    00f4495 View commit details
    Browse the repository at this point in the history

Commits on Jul 29, 2024

  1. final style core_test! test_jit_43 done!!! 

    Co-authored-by: Dougal Maclaurin <dougalm@google.com>
    @mattjj @dougalm
    mattjj and dougalm committed Jul 29, 2024
    Configuration menu
    Copy the full SHA
    e23b8b1 View commit details
    Browse the repository at this point in the history

Commits on Jul 30, 2024

  1. DynamicJaxprTrace.process_call should call to_jaxpr_tracer

    @mattjj
    mattjj committed Jul 30, 2024
    Configuration menu
    Copy the full SHA
    f8194e2 View commit details
    Browse the repository at this point in the history

  2. fix refs to DynamicJaxprTrace.instantiate_const

    @dougalm
    dougalm committed Jul 30, 2024
    Configuration menu
    Copy the full SHA
    9777c27 View commit details
    Browse the repository at this point in the history

  3. Add a `take_current_trace ctx manager to mostly replace find_cur_trac… 

    …e and avoid mixing up explicit and implicit traces
    @dougalm
    dougalm committed Jul 30, 2024
    Configuration menu
    Copy the full SHA
    39ad05c View commit details
    Browse the repository at this point in the history

  4. started working on shard_map, one test working

    @mattjj
    mattjj committed Jul 30, 2024
    Configuration menu
    Copy the full SHA
    5cfb8e1 View commit details
    Browse the repository at this point in the history

  5. Merge branch 'stackless' of github.com:google/jax into stackless

    @dougalm
    dougalm committed Jul 30, 2024
    Configuration menu
    Copy the full SHA
    7dcd2b2 View commit details
    Browse the repository at this point in the history

  6. Shard map tests 

    Co-authored-by: Matt Johnson <mattjj@google.com>
    @dougalm @mattjj
    dougalm and mattjj committed Jul 30, 2024
    Configuration menu
    Copy the full SHA
    7c8010a View commit details
    Browse the repository at this point in the history

Commits on Jul 31, 2024

  1. batching tests

    @dougalm
    dougalm committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    1172dd0 View commit details
    Browse the repository at this point in the history

Commits on Aug 1, 2024

  1. Avoid extending the axis mesh twice (todo: figure out a proper story … 

    …for mesh contexts and when they're implicit vs explicit)
    @dougalm
    dougalm committed Aug 1, 2024
    Configuration menu
    Copy the full SHA
    67a15ff View commit details
    Browse the repository at this point in the history

Commits on Aug 5, 2024

  1. Custom vmap tests

    @dougalm
    dougalm committed Aug 5, 2024
    Configuration menu
    Copy the full SHA
    9c6ca68 View commit details
    Browse the repository at this point in the history

  2. custom transpose tests

    @dougalm
    dougalm committed Aug 5, 2024
    Configuration menu
    Copy the full SHA
    ae0b448 View commit details
    Browse the repository at this point in the history

  3. more custom vjp/jvp rules

    @dougalm
    dougalm committed Aug 5, 2024
    Configuration menu
    Copy the full SHA
    b75ac75 View commit details
    Browse the repository at this point in the history

  4. more custom ad

    @dougalm
    dougalm committed Aug 5, 2024
    Configuration menu
    Copy the full SHA
    18af559 View commit details
    Browse the repository at this point in the history

  5. Avoid adding no_axis_name to env

    @dougalm
    dougalm committed Aug 5, 2024
    Configuration menu
    Copy the full SHA
    db38066 View commit details
    Browse the repository at this point in the history

Commits on Aug 6, 2024

  1. reinstate leak checker

    @dougalm
    dougalm committed Aug 6, 2024
    Configuration menu
    Copy the full SHA
    612cd17 View commit details
    Browse the repository at this point in the history

  2. avoid false positive tracer leak errors due to frame.constid_to_tracer

    @dougalm
    dougalm committed Aug 6, 2024
    Configuration menu
    Copy the full SHA
    450282c View commit details
    Browse the repository at this point in the history

  3. jet

    @dougalm
    dougalm committed Aug 6, 2024
    Configuration menu
    Copy the full SHA
    038217f View commit details
    Browse the repository at this point in the history

  4. sparsify tests

    @dougalm
    dougalm committed Aug 6, 2024
    Configuration menu
    Copy the full SHA
    778b3bb View commit details
    Browse the repository at this point in the history

  5. pmap tests

    @dougalm
    dougalm committed Aug 6, 2024
    Configuration menu
    Copy the full SHA
    b937243 View commit details
    Browse the repository at this point in the history

  6. more pmap

    @dougalm
    dougalm committed Aug 6, 2024
    Configuration menu
    Copy the full SHA
    3bf4a2a View commit details
    Browse the repository at this point in the history

Commits on Aug 7, 2024

  1. pallas call tests

    @dougalm
    dougalm committed Aug 7, 2024
    Configuration menu
    Copy the full SHA
    cf7da64 View commit details
    Browse the repository at this point in the history

  2. pmap tests

    @dougalm
    dougalm committed Aug 7, 2024
    Configuration menu
    Copy the full SHA
    3ab29f5 View commit details
    Browse the repository at this point in the history

  3. more pmap

    @dougalm
    dougalm committed Aug 7, 2024
    Configuration menu
    Copy the full SHA
    0136695 View commit details
    Browse the repository at this point in the history

Commits on Aug 8, 2024

  1. callback tests

    @dougalm
    dougalm committed Aug 8, 2024
    Configuration menu
    Copy the full SHA
    b7244b2 View commit details
    Browse the repository at this point in the history

  2. fix some pmap tests, 161 -> 126 pmap_test failures

    @mattjj
    mattjj committed Aug 8, 2024
    Configuration menu
    Copy the full SHA
    4a624d8 View commit details
    Browse the repository at this point in the history

  3. remove breakpoints (sorry)

    @mattjj
    mattjj committed Aug 8, 2024
    Configuration menu
    Copy the full SHA
    25e105b View commit details
    Browse the repository at this point in the history

Commits on Aug 9, 2024

  1. sparse transform infinite recursion

    @dougalm
    dougalm committed Aug 9, 2024
    Configuration menu
    Copy the full SHA
    9b428ae View commit details
    Browse the repository at this point in the history

  2. use correct tangent space dtype for zero in sparse array constructors

    @dougalm
    dougalm committed Aug 9, 2024
    Configuration menu
    Copy the full SHA
    babe549 View commit details
    Browse the repository at this point in the history

  3. Avoid cyclic refs so we don't have to run gc during trace leak checking

    @dougalm
    dougalm committed Aug 9, 2024
    Configuration menu
    Copy the full SHA
    9ed34db View commit details
    Browse the repository at this point in the history

Commits on Aug 10, 2024

  1. fix all non-eager pmap tests

    @mattjj
    mattjj committed Aug 10, 2024
    Configuration menu
    Copy the full SHA
    2f8c03f View commit details
    Browse the repository at this point in the history

  2. fix pmap bind logic, down to 24 failing pmap tests

    @mattjj
    mattjj committed Aug 10, 2024
    Configuration menu
    Copy the full SHA
    1320052 View commit details
    Browse the repository at this point in the history

  3. fix eager pmap axis_index

    @mattjj
    mattjj committed Aug 10, 2024
    Configuration menu
    Copy the full SHA
    44087f5 View commit details
    Browse the repository at this point in the history

  4. fix nested eager pmap

    @mattjj
    mattjj committed Aug 10, 2024
    Configuration menu
    Copy the full SHA
    af30a2a View commit details
    Browse the repository at this point in the history

  5. more eager pmap fixes

    @mattjj
    mattjj committed Aug 10, 2024
    Configuration menu
    Copy the full SHA
    bbad195 View commit details
    Browse the repository at this point in the history

  6. fix last of the eager pmap tests

    @mattjj
    mattjj committed Aug 10, 2024
    Configuration menu
    Copy the full SHA
    59ea325 View commit details
    Browse the repository at this point in the history

Commits on Aug 11, 2024

  1. skip some for_loop tests that were timing out

    @dougalm
    dougalm committed Aug 11, 2024
    Configuration menu
    Copy the full SHA
    f382187 View commit details
    Browse the repository at this point in the history

Commits on Aug 12, 2024

  1. Extra short-circuit for not-implemented batching rules

    @dougalm
    dougalm committed Aug 12, 2024
    Configuration menu
    Copy the full SHA
    5ab7e3c View commit details
    Browse the repository at this point in the history

  2. Add a special case for DimExpr

    @dougalm
    dougalm committed Aug 12, 2024
    Configuration menu
    Copy the full SHA
    28e9383 View commit details
    Browse the repository at this point in the history

  3. set current trace in pjit dynamic shapes staging rule

    @dougalm
    dougalm committed Aug 12, 2024
    Configuration menu
    Copy the full SHA
    e83bb5b View commit details
    Browse the repository at this point in the history

  4. more dynamic_api_test

    @dougalm
    dougalm committed Aug 12, 2024
    Configuration menu
    Copy the full SHA
    4bf0fe1 View commit details
    Browse the repository at this point in the history

  5. reset trace state between tests

    @dougalm
    dougalm committed Aug 12, 2024
    Configuration menu
    Copy the full SHA
    5a40ee5 View commit details
    Browse the repository at this point in the history

Commits on Aug 13, 2024

  1. pallas tests

    @dougalm
    dougalm committed Aug 13, 2024
    Configuration menu
    Copy the full SHA
    0cf761c View commit details
    Browse the repository at this point in the history

  2. skipping infeed tests

    @dougalm
    dougalm committed Aug 13, 2024
    Configuration menu
    Copy the full SHA
    f52aa26 View commit details
    Browse the repository at this point in the history

  3. xla_computation

    @dougalm
    dougalm committed Aug 13, 2024
    Configuration menu
    Copy the full SHA
    3120391 View commit details
    Browse the repository at this point in the history

  4. dead trace invalidation

    @dougalm
    dougalm committed Aug 13, 2024
    Configuration menu
    Copy the full SHA
    26e5fb5 View commit details
    Browse the repository at this point in the history

Commits on Aug 14, 2024

  1. Reset name stack when tracing out a jaxpr

    @dougalm
    dougalm committed Aug 14, 2024
    Configuration menu
    Copy the full SHA
    3442473 View commit details
    Browse the repository at this point in the history

  2. Add the eval context back during callbacks

    @dougalm
    dougalm committed Aug 14, 2024
    Configuration menu
    Copy the full SHA
    ccf9877 View commit details
    Browse the repository at this point in the history

  3. Custom vjp/jvp closing over stuff

    @dougalm
    dougalm committed Aug 14, 2024
    Configuration menu
    Copy the full SHA
    296b299 View commit details
    Browse the repository at this point in the history

  4. pallas tests

    @dougalm
    dougalm committed Aug 14, 2024
    Configuration menu
    Copy the full SHA
    fa29f36 View commit details
    Browse the repository at this point in the history

  5. Sparse transform test

    @dougalm
    dougalm committed Aug 14, 2024
    Configuration menu
    Copy the full SHA
    bdd18c4 View commit details
    Browse the repository at this point in the history

  6. mutable arrays

    @dougalm
    dougalm committed Aug 14, 2024
    Configuration menu
    Copy the full SHA
    99d905f View commit details
    Browse the repository at this point in the history

  7. test fix

    @dougalm
    dougalm committed Aug 14, 2024
    Configuration menu
    Copy the full SHA
    30376dd View commit details
    Browse the repository at this point in the history

  8. effects test

    @dougalm
    dougalm committed Aug 14, 2024
    Configuration menu
    Copy the full SHA
    77c09f1 View commit details
    Browse the repository at this point in the history

Commits on Aug 26, 2024

  1. Check whether a tracer's trace is invalidated

    @dougalm
    dougalm committed Aug 26, 2024
    Configuration menu
    Copy the full SHA
    b2edd22 View commit details
    Browse the repository at this point in the history

  2. tweak escaped tracer tests

    @dougalm
    dougalm committed Aug 26, 2024
    Configuration menu
    Copy the full SHA
    aeae13d View commit details
    Browse the repository at this point in the history

  3. Remove CustomJVPException (only relevant for post_process_call)

    @dougalm
    dougalm committed Aug 26, 2024
    Configuration menu
    Copy the full SHA
    16af376 View commit details
    Browse the repository at this point in the history

  4. more tweaks to escaped tracer tests

    @dougalm
    dougalm committed Aug 26, 2024
    Configuration menu
    Copy the full SHA
    10f71ca View commit details
    Browse the repository at this point in the history

Commits on Aug 28, 2024

  1. Use primal-type-to-tangent-type mapping when creating symbolic zeros … 

    …from primals
    @dougalm
    dougalm committed Aug 28, 2024
    Configuration menu
    Copy the full SHA
    9015af1 View commit details
    Browse the repository at this point in the history

  2. wip

    @dougalm
    dougalm committed Aug 28, 2024
    Configuration menu
    Copy the full SHA
    6b9b191 View commit details
    Browse the repository at this point in the history

  3. Merge branch 'main' into stackless

    @dougalm
    dougalm committed Aug 28, 2024
    Configuration menu
    Copy the full SHA
    f4200ba View commit details
    Browse the repository at this point in the history

  4. fix aval error

    @dougalm
    dougalm committed Aug 28, 2024
    Configuration menu
    Copy the full SHA
    c1f871a View commit details
    Browse the repository at this point in the history

  5. update remat opt

    @dougalm
    dougalm committed Aug 28, 2024
    Configuration menu
    Copy the full SHA
    4944e2f View commit details
    Browse the repository at this point in the history

Commits on Aug 29, 2024

  1. Use correct tangent space avals in custom jvp thunk

    @dougalm
    dougalm committed Aug 29, 2024
    Configuration menu
    Copy the full SHA
    2d854f5 View commit details
    Browse the repository at this point in the history

  2. move promotion-to-float outside of logaddexp primitive to avoid float… 

    …0 tangents
    @dougalm
    dougalm committed Aug 29, 2024
    Configuration menu
    Copy the full SHA
    5d427aa View commit details
    Browse the repository at this point in the history

Commits on Aug 30, 2024

  1. update attrs to use take_current_trace

    @dougalm
    dougalm committed Aug 30, 2024
    Configuration menu
    Copy the full SHA
    df488b0 View commit details
    Browse the repository at this point in the history

Commits on Sep 3, 2024

  1. Organize tracing context, hoping to get caching working again

    @dougalm
    dougalm committed Sep 3, 2024
    Configuration menu
    Copy the full SHA
    3a56e67 View commit details
    Browse the repository at this point in the history

  2. Add axis size to cache key

    @dougalm
    dougalm committed Sep 3, 2024
    Configuration menu
    Copy the full SHA
    eb64ddc View commit details
    Browse the repository at this point in the history

  3. Use weak refs to traces in cache key

    @dougalm
    dougalm committed Sep 3, 2024
    Configuration menu
    Copy the full SHA
    4ede817 View commit details
    Browse the repository at this point in the history

  4. Fix

    @dougalm
    dougalm committed Sep 3, 2024
    Configuration menu
    Copy the full SHA
    6ea0c68 View commit details
    Browse the repository at this point in the history

  5. fixes

    @dougalm
    dougalm committed Sep 3, 2024
    Configuration menu
    Copy the full SHA
    7354526 View commit details
    Browse the repository at this point in the history

  6. set eval trace for key reuse checker

    @dougalm
    dougalm committed Sep 3, 2024
    Configuration menu
    Copy the full SHA
    1a32b9a View commit details
    Browse the repository at this point in the history

  7. set trace for dimension_as_value

    @dougalm
    dougalm committed Sep 3, 2024
    Configuration menu
    Copy the full SHA
    67c38cf View commit details
    Browse the repository at this point in the history

  8. fix pallas primitive custom bind logic

    @dougalm
    dougalm committed Sep 3, 2024
    Configuration menu
    Copy the full SHA
    420f40d View commit details
    Browse the repository at this point in the history

  9. Fix attribute errors

    @dougalm
    dougalm committed Sep 3, 2024
    Configuration menu
    Copy the full SHA
    46af63b View commit details
    Browse the repository at this point in the history

Commits on Sep 4, 2024

  1. Add back partial_eval._memoize with a patently incorrect new implemen… 

    …tation
    @dougalm
    dougalm committed Sep 4, 2024
    Configuration menu
    Copy the full SHA
    4167de0 View commit details
    Browse the repository at this point in the history

  2. Use an empty context when forcing jaxpr thunks

    @dougalm
    dougalm committed Sep 4, 2024
    Configuration menu
    Copy the full SHA
    12cd781 View commit details
    Browse the repository at this point in the history

  3. fix invalidation negation

    @dougalm
    dougalm committed Sep 4, 2024
    Configuration menu
    Copy the full SHA
    2566b6a View commit details
    Browse the repository at this point in the history

  4. unexpected tracer error messages

    @dougalm
    dougalm committed Sep 4, 2024
    Configuration menu
    Copy the full SHA
    acd4bb7 View commit details
    Browse the repository at this point in the history

  5. small fixes

    @dougalm
    dougalm committed Sep 4, 2024
    Configuration menu
    Copy the full SHA
    49fc277 View commit details
    Browse the repository at this point in the history

  6. tangent dtype conversion in test

    @dougalm
    dougalm committed Sep 4, 2024
    Configuration menu
    Copy the full SHA
    5351c62 View commit details
    Browse the repository at this point in the history

  7. Add back a deleted check in bind

    @dougalm
    dougalm committed Sep 4, 2024
    Configuration menu
    Copy the full SHA
    34e1fe9 View commit details
    Browse the repository at this point in the history

  8. revert - need to handle function arguments

    @dougalm
    dougalm committed Sep 4, 2024
    Configuration menu
    Copy the full SHA
    7159909 View commit details
    Browse the repository at this point in the history

  9. don't tempt fate by creating jvp tracers with symbolic zero tangents … 

    …too readily
    @dougalm
    dougalm committed Sep 4, 2024
    Configuration menu
    Copy the full SHA
    e89afc8 View commit details
    Browse the repository at this point in the history

  10. more float0

    @dougalm
    dougalm committed Sep 4, 2024
    Configuration menu
    Copy the full SHA
    64d10d7 View commit details
    Browse the repository at this point in the history

Commits on Sep 5, 2024

  1. skip a test

    @dougalm
    dougalm committed Sep 5, 2024
    Configuration menu
    Copy the full SHA
    2e96464 View commit details
    Browse the repository at this point in the history

  2. Merge branch 'main' into stackless

    @dougalm
    dougalm committed Sep 5, 2024
    Configuration menu
    Copy the full SHA
    b237bd2 View commit details
    Browse the repository at this point in the history

  3. Add back _convert_element_type custom bind

    @dougalm
    dougalm committed Sep 5, 2024
    Configuration menu
    Copy the full SHA
    f26c578 View commit details
    Browse the repository at this point in the history

  4. deleted xmap cruft

    @dougalm
    dougalm committed Sep 5, 2024
    Configuration menu
    Copy the full SHA
    b9f1f23 View commit details
    Browse the repository at this point in the history

  5. more deletion

    @dougalm
    dougalm committed Sep 5, 2024
    Configuration menu
    Copy the full SHA
    69cf62f View commit details
    Browse the repository at this point in the history

  6. delete whitespace

    @dougalm
    dougalm committed Sep 5, 2024
    Configuration menu
    Copy the full SHA
    2abe624 View commit details
    Browse the repository at this point in the history

Commits on Sep 6, 2024

  1. Fix pytype errors

    @dougalm
    dougalm committed Sep 6, 2024
    Configuration menu
    Copy the full SHA
    d0a28a0 View commit details
    Browse the repository at this point in the history

  2. Put TraceTag in core

    @dougalm
    dougalm committed Sep 6, 2024
    Configuration menu
    Copy the full SHA
    500cb43 View commit details
    Browse the repository at this point in the history

  3. fix

    @dougalm
    dougalm committed Sep 6, 2024
    Configuration menu
    Copy the full SHA
    0155827 View commit details
    Browse the repository at this point in the history

  4. more

    @dougalm
    dougalm committed Sep 6, 2024
    Configuration menu
    Copy the full SHA
    2d03ea6 View commit details
    Browse the repository at this point in the history

  5. process_custom_jvp for rewrite trace

    @dougalm
    dougalm committed Sep 6, 2024
    Configuration menu
    Copy the full SHA
    75be2e3 View commit details
    Browse the repository at this point in the history

  6. shard map rewrite custom_vjp

    @dougalm
    dougalm committed Sep 6, 2024
    Configuration menu
    Copy the full SHA
    f22d8c0 View commit details
    Browse the repository at this point in the history

  7. shard_map batching

    @dougalm
    dougalm committed Sep 6, 2024
    Configuration menu
    Copy the full SHA
    9a8beab View commit details
    Browse the repository at this point in the history

  8. more shard map rewrites

    @dougalm
    dougalm committed Sep 6, 2024
    Configuration menu
    Copy the full SHA
    d242e81 View commit details
    Browse the repository at this point in the history

  9. even more rewrite

    @dougalm
    dougalm committed Sep 6, 2024
    Configuration menu
    Copy the full SHA
    6926716 View commit details
    Browse the repository at this point in the history

Commits on Sep 7, 2024

  1. axis env

    @dougalm
    dougalm committed Sep 7, 2024
    Configuration menu
    Copy the full SHA
    e32c5b9 View commit details
    Browse the repository at this point in the history

Commits on Sep 9, 2024

  1. axis index custom bind

    @dougalm
    dougalm committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    7e5a6b4 View commit details
    Browse the repository at this point in the history

  2. ptype errors

    @dougalm
    dougalm committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    4f2fc87 View commit details
    Browse the repository at this point in the history

  3. lint errors

    @dougalm
    dougalm committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    ed32420 View commit details
    Browse the repository at this point in the history

  4. fix to lint fix

    @dougalm
    dougalm committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    9dee9f1 View commit details
    Browse the repository at this point in the history

  5. revert core test monkey patch

    @dougalm
    dougalm committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    49e5c57 View commit details
    Browse the repository at this point in the history

  6. remove breakpoints

    @dougalm
    dougalm committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    ceefd62 View commit details
    Browse the repository at this point in the history

  7. tweak docstring for test

    @dougalm
    dougalm committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    a591af0 View commit details
    Browse the repository at this point in the history

  8. fix pytype error

    @dougalm
    dougalm committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    aa5bcbd View commit details
    Browse the repository at this point in the history

  9. keep mypy happy by removing type:ignore ?

    @dougalm
    dougalm committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    bacd510 View commit details
    Browse the repository at this point in the history

  10. Merge branch 'main' into stackless

    @dougalm
    dougalm committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    ee130dd View commit details
    Browse the repository at this point in the history

Commits on Sep 10, 2024

  1. helper method for overriding bind_with_trace

    @dougalm
    dougalm committed Sep 10, 2024
    Configuration menu
    Copy the full SHA
    8dd89ac View commit details
    Browse the repository at this point in the history

  2. comment out jax2tf stuff that's not ready yet to satisfy mypy

    @dougalm
    dougalm committed Sep 10, 2024
    Configuration menu
    Copy the full SHA
    873550a View commit details
    Browse the repository at this point in the history

  3. reword a line to try to make mypy happy?

    @dougalm
    dougalm committed Sep 10, 2024
    Configuration menu
    Copy the full SHA
    640ea0f View commit details
    Browse the repository at this point in the history

  4. check for shapedarray

    @dougalm
    dougalm committed Sep 10, 2024
    Configuration menu
    Copy the full SHA
    788c37a View commit details
    Browse the repository at this point in the history

  5. mypy

    @dougalm
    dougalm committed Sep 10, 2024
    Configuration menu
    Copy the full SHA
    c83e492 View commit details
    Browse the repository at this point in the history

  6. jax2tf

    @dougalm
    dougalm committed Sep 10, 2024
    Configuration menu
    Copy the full SHA
    758879d View commit details
    Browse the repository at this point in the history

  7. Merge branch 'main' into stackless

    @dougalm
    dougalm committed Sep 10, 2024
    Configuration menu
    Copy the full SHA
    fe49a2c View commit details
    Browse the repository at this point in the history

  8. fix all_to_all batching rule

    @dougalm
    dougalm committed Sep 10, 2024
    Configuration menu
    Copy the full SHA
    96d751b View commit details
    Browse the repository at this point in the history

Commits on Sep 11, 2024

  1. tf tracing

    @dougalm
    dougalm committed Sep 11, 2024
    Configuration menu
    Copy the full SHA
    612e941 View commit details
    Browse the repository at this point in the history

  2. missed one

    @dougalm
    dougalm committed Sep 11, 2024
    Configuration menu
    Copy the full SHA
    d3bf703 View commit details
    Browse the repository at this point in the history

Commits on Sep 12, 2024

  1. trace tag hash hack

    @dougalm
    dougalm committed Sep 12, 2024
    Configuration menu
    Copy the full SHA
    0f52c8a View commit details
    Browse the repository at this point in the history

Commits on Sep 13, 2024

  1. Merge branch 'main' into stackless

    @dougalm
    dougalm committed Sep 13, 2024
    Configuration menu
    Copy the full SHA
    2af9167 View commit details
    Browse the repository at this point in the history

  2. lint

    @dougalm
    dougalm committed Sep 13, 2024
    Configuration menu
    Copy the full SHA
    6b674c2 View commit details
    Browse the repository at this point in the history

  3. skip test

    @dougalm
    dougalm committed Sep 13, 2024
    Configuration menu
    Copy the full SHA
    6e8828d View commit details
    Browse the repository at this point in the history

  4. Add leak checker to batch tracing

    @dougalm
    dougalm committed Sep 13, 2024
    Configuration menu
    Copy the full SHA
    885b3a0 View commit details
    Browse the repository at this point in the history

  5. oof leak checker false positives. this work?

    @dougalm
    dougalm committed Sep 13, 2024
    Configuration menu
    Copy the full SHA
    874fa26 View commit details
    Browse the repository at this point in the history

  6. skip const-forwarding test

    @dougalm
    dougalm committed Sep 13, 2024
    Configuration menu
    Copy the full SHA
    3676cc9 View commit details
    Browse the repository at this point in the history

  7. tweak a test

    @dougalm
    dougalm committed Sep 13, 2024
    Configuration menu
    Copy the full SHA
    fe3e6b4 View commit details
    Browse the repository at this point in the history

Commits on Sep 14, 2024

  1. use tangent type in custom lin

    @dougalm
    dougalm committed Sep 14, 2024
    Configuration menu
    Copy the full SHA
    b7cfa55 View commit details
    Browse the repository at this point in the history

Commits on Sep 16, 2024

  1. Avoid blowing away custom jvp rule during partial eval

    @dougalm
    dougalm committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    8f2c429 View commit details
    Browse the repository at this point in the history

  2. Merge branch 'main' into stackless

    @dougalm
    dougalm committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    65e0922 View commit details
    Browse the repository at this point in the history

  3. unused import

    @dougalm
    dougalm committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    bef2066 View commit details
    Browse the repository at this point in the history

  4. Rename at_least_vspace -> to_tangent_aval, from_value -> `from_… 

    …primal_value`.
    @dougalm
    dougalm committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    2604e50 View commit details
    Browse the repository at this point in the history

  5. Fix return type

    @dougalm
    dougalm committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    227e50f View commit details
    Browse the repository at this point in the history

  6. Use tangent types in appropriate places

    @dougalm
    dougalm committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    62b3ac1 View commit details
    Browse the repository at this point in the history

  7. fix tests (remove replace/recast float0)

    @dougalm
    dougalm committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    be0cc23 View commit details
    Browse the repository at this point in the history

  8. missed one

    @dougalm
    dougalm committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    00abb18 View commit details
    Browse the repository at this point in the history

  9. fix another dtype bug in tests

    @dougalm
    dougalm committed Sep 16, 2024
    Configuration menu
    Copy the full SHA
    c03df59 View commit details
    Browse the repository at this point in the history

Commits on Sep 17, 2024

  1. Maybe we can put this check back (but really I just want to trigger c… 

    …opybara
    @dougalm
    dougalm committed Sep 17, 2024
    Configuration menu
    Copy the full SHA
    99760a0 View commit details
    Browse the repository at this point in the history

  2. revert

    @dougalm
    dougalm committed Sep 17, 2024
    Configuration menu
    Copy the full SHA
    a5c4459 View commit details
    Browse the repository at this point in the history

  3. Merge branch 'tangent-avals' into stackless

    @dougalm
    dougalm committed Sep 17, 2024
    Configuration menu
    Copy the full SHA
    7d5c660 View commit details
    Browse the repository at this point in the history

  4. docstring tweak for test

    @dougalm
    dougalm committed Sep 17, 2024
    Configuration menu
    Copy the full SHA
    b8c17ca View commit details
    Browse the repository at this point in the history

  5. tweak xla_metadata_test to avoid using impl path directly

    @dougalm
    dougalm committed Sep 17, 2024
    Configuration menu
    Copy the full SHA
    c80594f View commit details
    Browse the repository at this point in the history

Commits on Sep 18, 2024

  1. Add back NamedAxisEffect

    @dougalm
    dougalm committed Sep 18, 2024
    Configuration menu
    Copy the full SHA
    a4117e2 View commit details
    Browse the repository at this point in the history

  2. fix

    @dougalm
    dougalm committed Sep 18, 2024
    Configuration menu
    Copy the full SHA
    b17e5af View commit details
    Browse the repository at this point in the history

  3. named axis reversion fixes

    @dougalm
    dougalm committed Sep 18, 2024
    Configuration menu
    Copy the full SHA
    1b8d743 View commit details
    Browse the repository at this point in the history

  4. more fixes

    @dougalm
    dougalm committed Sep 18, 2024
    Configuration menu
    Copy the full SHA
    13139fe View commit details
    Browse the repository at this point in the history

  5. Merge branch 'main' into stackless

    @dougalm
    dougalm committed Sep 18, 2024
    Configuration menu
    Copy the full SHA
    8b72dae View commit details
    Browse the repository at this point in the history

Commits on Sep 19, 2024

  1. missed a filter_named_axis_effects

    @dougalm
    dougalm committed Sep 19, 2024
    Configuration menu
    Copy the full SHA
    f9dd1e6 View commit details
    Browse the repository at this point in the history

  2. remove dead code

    @dougalm
    dougalm committed Sep 19, 2024
    Configuration menu
    Copy the full SHA
    a131df7 View commit details
    Browse the repository at this point in the history

  3. Merge branch 'main' into stackless

    @dougalm
    dougalm committed Sep 19, 2024
    Configuration menu
    Copy the full SHA
    cc5c9d6 View commit details
    Browse the repository at this point in the history

  4. Merge branch 'main' into stackless

    @dougalm
    dougalm committed Sep 19, 2024
    Configuration menu
    Copy the full SHA
    ca6e762 View commit details
    Browse the repository at this point in the history

Commits on Sep 20, 2024

  1. fix bad merge

    @dougalm
    dougalm committed Sep 20, 2024
    Configuration menu
    Copy the full SHA
    f9b3c49 View commit details
    Browse the repository at this point in the history

  2. rearrange to shrink diff a bit

    @dougalm
    dougalm committed Sep 20, 2024
    Configuration menu
    Copy the full SHA
    13eb6e2 View commit details
    Browse the repository at this point in the history

  3. more diff tweaks

    @dougalm
    dougalm committed Sep 20, 2024
    Configuration menu
    Copy the full SHA
    d5f6166 View commit details
    Browse the repository at this point in the history

  4. more minor stuff

    @dougalm
    dougalm committed Sep 20, 2024
    Configuration menu
    Copy the full SHA
    efd6c21 View commit details
    Browse the repository at this point in the history

  5. fix

    @dougalm
    dougalm committed Sep 20, 2024
    Configuration menu
    Copy the full SHA
    3d3dd93 View commit details
    Browse the repository at this point in the history

Commits on Sep 21, 2024

  1. Update batching axis size as you go under a shard map

    @dougalm
    dougalm committed Sep 21, 2024
    Configuration menu
    Copy the full SHA
    e867e66 View commit details
    Browse the repository at this point in the history

  2. Merge branch 'main' into stackless

    @dougalm
    dougalm committed Sep 21, 2024
    Configuration menu
    Copy the full SHA
    a32ab6a View commit details
    Browse the repository at this point in the history

  3. more batching cleanup

    @dougalm
    dougalm committed Sep 21, 2024
    Configuration menu
    Copy the full SHA
    3dffee1 View commit details
    Browse the repository at this point in the history

  4. more batching

    @dougalm
    dougalm committed Sep 21, 2024
    Configuration menu
    Copy the full SHA
    2df3c20 View commit details
    Browse the repository at this point in the history

  5. more batching

    @dougalm
    dougalm committed Sep 21, 2024
    Configuration menu
    Copy the full SHA
    d8303a4 View commit details
    Browse the repository at this point in the history

  6. fix

    @dougalm
    dougalm committed Sep 21, 2024
    Configuration menu
    Copy the full SHA
    abf4d98 View commit details
    Browse the repository at this point in the history

  7. jax2tf fix

    @dougalm
    dougalm committed Sep 21, 2024
    Configuration menu
    Copy the full SHA
    213733d View commit details
    Browse the repository at this point in the history

Commits on Sep 23, 2024

  1. Merge branch 'main' into stackless

    @dougalm
    dougalm committed Sep 23, 2024
    Configuration menu
    Copy the full SHA
    44709e0 View commit details
    Browse the repository at this point in the history

  2. small backwards compat shims

    @dougalm
    dougalm committed Sep 23, 2024
    Configuration menu
    Copy the full SHA
    74d21ec View commit details
    Browse the repository at this point in the history

  3. fix merge

    @dougalm
    dougalm committed Sep 23, 2024
    Configuration menu
    Copy the full SHA
    b0533c3 View commit details
    Browse the repository at this point in the history

Commits on Sep 24, 2024

  1. Merge branch 'main' into stackless

    @dougalm
    dougalm committed Sep 24, 2024
    Configuration menu
    Copy the full SHA
    6d068f2 View commit details
    Browse the repository at this point in the history

Commits on Sep 25, 2024

  1. fix bad merge

    @dougalm
    dougalm committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    63523c7 View commit details
    Browse the repository at this point in the history

  2. more fix

    @dougalm
    dougalm committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    5df8920 View commit details
    Browse the repository at this point in the history

  3. implement unsafe trace-querying APIs

    @dougalm
    dougalm committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    0b24498 View commit details
    Browse the repository at this point in the history

  4. lint

    @dougalm
    dougalm committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    23d3e01 View commit details
    Browse the repository at this point in the history

  5. Merge branch 'main' into stackless

    @dougalm
    dougalm committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    6b90f6d View commit details
    Browse the repository at this point in the history