jax
e854f165 - Allow `P.UNCONSTRAINED` in out_shardings at top level jit. This is required for sharding in types to work properly when out_avals contain UNCONSTRAINED specs.

Commit

1 year ago

Allow `P.UNCONSTRAINED` in out_shardings at top level jit. This is required for sharding in types to work properly when out_avals contain UNCONSTRAINED specs. This also simplifies the `impl` rule of `sharding_cast`. PiperOrigin-RevId: 707349491

References

#25524 - Allow `P.UNCONSTRAINED` in out_shardings at top level jit. This is required for sharding in types to work properly when out_avals contain UNCONSTRAINED specs.

#187 - CI: 12/19/24 upstream sync

Author

yashk2810

Committer

Google-ML-Automation

Parents

b56dc631

jax e854f165 - Allow `P.UNCONSTRAINED` in out_shardings at top level jit. This is required for sharding in types to work properly when out_avals contain UNCONSTRAINED specs.

jax
e854f165 - Allow `P.UNCONSTRAINED` in out_shardings at top level jit. This is required for sharding in types to work properly when out_avals contain UNCONSTRAINED specs.