Copybara import of the project:
Copybara import of the project:

--
740cc91382efe9933d4f5eea5f975cf3a6aa8d34 by George Necula <gcnecula@gmail.com>:

[masking] Ensure that Poly operations are sound

We want to ensure that an operation on polynomials throws
an explicit error when the result of the operation cannot
be soundly represented as a polynomial.

This change implements polynomials ranging over non-negative
shape variables. (It may be a good idea to restrict the range
to positive values only.)

Formally, an operation `op` on polynomial `p` either
throws an error `UndefinedPoly` or produces a polynomial
`res` such that `op(Val(p)) = Val(res)`, for any `Val`, a
value with non-negative values for shape variables.

Prior to this there were several sources of unsoundness:

 * The comparison `p1 = p2` would return True when the
   polynomials are identical (sound), and False when they
   are not identical (not sound)
 * The comparison `p >= 0` would always return True. This
   is unsound if `p` has monomial with negative coefficients.

--
df081393ebded602c9007361f5d988945ea2e95d by Matthew Johnson <mattjj@google.com>:

update concatenate shape rule for sound mask test

--
799325bbba352842a56ca225f02b6eb970f45712 by Matthew Johnson <mattjj@google.com>:

revise lax broadcasting to work with sound masking

--
8bd74a245a878dbd1aff4d5de491fcdd3f889dd3 by Matthew Johnson <mattjj@google.com>:

adapt broadcast_in_dim shape rule mask soundness

--
9a0207ab30d3b08476edd032f945984e2a3b8381 by Matthew Johnson <mattjj@google.com>:

adapt masking rev tests for soundness

--
3bf3e352bc45cd6d61bc322501502595b42b3da5 by George Necula <gcnecula@gmail.com>:

Disable the remaining failing tests, remove jax2tf documentation

--
64022922bd18eacfcd9befd2e1810a80512f85ef by George Necula <gcnecula@gmail.com>:

Fix float indices in lax.slice

--
06f60f46b628cbf6c43dfb2ba013dc20c07b938f by George Necula <gcnecula@gmail.com>:

Disable more tests

PiperOrigin-RevId: 343095255

google/flax,google/flax