Difficulties converting with jax2tf and polymorphism

[thekevinscott]

Hello,

I'm new to Jax and struggling to convert a Jax model (MAXIM) to Tensorflow.js.

Converting the model with a variable batch dimension and fixed image dimensions works well:

tfjs.converters.convert_jax(
    jax_model.apply,
    params, 
    model_dir='./tmp', 
    input_signatures=[tf.TensorSpec([None, 256, 256, 3], tf.float32)],
    polymorphic_shapes=["(b, 256, 256, 3)"],    
)

However, for certain models fixed input sizes result in a noticeable patching effect. Therefore, I'm trying to update the model to support dynamic image inputs.

Changing the conversion code to something like:

tfjs.converters.convert_jax(
    jax_model.apply,
    params, 
    model_dir='./tmp', 
    input_signatures=[tf.TensorSpec([None, None, 256, 3], tf.float32)],
    polymorphic_shapes=["(b, h, 256, 3)"],    
)

Fails.

I've refactored the library's Jax code a bit in order to avoid the use of einops, so I'm not entirely sure how to share a runnable code snippet, but I can share the errors I'm seeing.

The code fails at:

    reshaped_padded = np.reshape(x_pad, new_shape)

With the values of x_pad and new_shape being:

x_pad:
Traced<ShapedArray(float32[b,h,512,32])>with<TensorFlowTrace(level=0/1)> with
  val = <tf.Tensor 'jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/pjit__pad_/PadV2_7:0' shape=(None, None, 512, 32) dtype=float32>
  _aval = ShapedArray(float32[b,h,512,32])

new_shape:
[b, floordiv(h, floordiv(h, 16)), floordiv(h, 16), 16, 32, 32]

This code fails an assertion here. When I dig in, b is indeed 0 (which fails the assertion) but I don't understand why. I assume it has to do with the use of the variable h and the floordivs but I'm hitting the limits of my knowledge wrt jax and its internals.

Can anyone point me in the right direction? I'm happy to provide more information if it's helpful.

---

UPDATE 1

Ok, I was able to whittle it down to the offending code. But I don't know why this fails or where in the docs to look for more information. Would appreciate any pointers!

This fails:

from jax.experimental.jax2tf.shape_poly import _DimExpr

h = _DimExpr.from_var('h')
r = h // (h // 16)
r.bounds()

Because this code returns a 0 in the tuple response:

(h // 16).bounds()

(0, inf)

---

UPDATE 2

So I continue digging!

from jax.experimental.jax2tf.shape_poly import _DimAtom, _DimExpr

h = _DimExpr.from_var('h')
r = h // 16

for mon, coeff in r.monomials():
    for a, exp in mon.items():
        for opnd in a.operands:
            print('Bounds for DimExpr', opnd.bounds())
        print('Bounds for DimAtom', a.bounds())

Produces:

Bounds for DimExpr (1, inf)
Bounds for DimExpr (16, 16)
Bounds for DimAtom (0, inf)

And what appears to be happening is that, on this line, the result of math.floor(1 / 16) becomes 0 which results in the error.

I further presume that the value of h gets cast as (1, inf) as the range of possible values.

So I think that leads me to ask: is there a way, when specifying dynamic variables, to specify a minimum value that a variable can be? E.g., is there a way to specify in the polymorphism settings that a value should be in the range of (64, inf)?

---

UPDATE 3

I learned from this that you can specify a dimension variable to be a multiple:

polymorphic_shapes=["(2*b, ...)"]

Which is great and exactly what I was looking for.

That leads me to a new error. The model that I'm converting has an Upsample layer that attempts to cast the height and width to concrete values. I'm trying to make this work with dynamic values.

I've tried the following:

    print(int(h * self.ratio))

yields:

jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(float32[], weak_type=True)>with<TensorFlowTrace(level=0/1)> with
  val = <tf.Tensor 'jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/pjit_fn_/Mul_1:0' shape=() dtype=float32>
  _aval = ShapedArray(float32[], weak_type=True)
The problem arose with the `int` function. If trying to convert the data type of a value, try using `x.astype(int)` or `jnp.array(x, int)` instead.

Fair enough. So I then try:

    print((h * self.ratio).astype(int))

Which yields:

Traced<ShapedArray(int32[])>with<TensorFlowTrace(level=0/1)> with
  val = <tf.Tensor 'jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Cast:0' shape=() dtype=int32>
  _aval = ShapedArray(int32[])

However, when I try to pass that to jax.image.resize:

    x = jax.image.resize(
        x,
        shape=(n, h, w, c),
        method="bilinear",

I get:

  File "/usr/local/lib/python3.8/dist-packages/jax/_src/image/scale.py", line 344, in resize
    return _resize(image, core.canonicalize_shape(shape), method, antialias,
TypeError: Shapes must be 1D sequences of concrete values of integer type, got (b, Traced<ShapedArray(int32[])>with<TensorFlowTrace(level=0/1)> with
  val = <tf.Tensor 'jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Cast_1:0' shape=() dtype=int32>
  _aval = ShapedArray(int32[]), 128, 32).

The readme indicates that this error:

means that one operation that computes shape parameters is using non-integer arguments, e.g., np.ndarray, that get implicitly converted to JAX arrays. The solution is to avoid np.array, float, or JAX arrays in operations whose results are used as shapes, e.g., instead of np.arange(n) * x.shape[0] write [i * x.shape[0] for i in range(n)].

Does this mean that jax.image.resize only accepts shape as a tuple of integers? If so, how does one write code that can support dynamic image sizes?

[thekevinscott]

This discussion is getting a bit long, so I've carved out the latest question into its own discussion.

The thrust of the original question was, is there a way to specify that a dynamic shape has a range? The answer is, sort of, you can specify that a dynamic variable can be a multiple:

polymorphic_shapes=["(2*b, ...)"]