Shape Polymorphism, with an image downscale

[thekevinscott]

Hi, I'd opened a prior discussion that got a bit long, so I'm restarting with a more specific question.

I want to export a model with Jax2tf that has polymorphic shapes and an image resizing. Specifically, I want to use a float instead of an integer, to result in a downscaled image.

Here's a sample (largely copied from here):

from jax.experimental import jax2tf

polymorphic_image_shape = jax2tf.shape_poly.PolyShape(
  1, 'height', 'width', 3
)
def apply_fn(image):
  b,h,w,c = image.shape
  h = h / 2
  return jax.image.resize(image, (b,h,w,c), method="linear")

apply_fn_jitted = jax.jit(apply_fn, backend='cpu')
apply_fn_tf = tf.function(
    jax2tf.convert(
        apply_fn_jitted,
        with_gradient=False,
        enable_xla=True,
        polymorphic_shapes=[polymorphic_image_shape]),
    autograph=False,
    jit_compile=True)

image_f = np.random.uniform(size=(1, 4, 4, 3))
apply_fn_tf(image_f)

This fails with:

TypeError: Shapes must be 1D sequences of concrete values of integer type, got (1, Traced<ShapedArray(float32[], weak_type=True)>with<TensorFlowTrace(level=0/2)> with
  val = <tf.Tensor 'jax2tf_apply_fn_/jit_apply_fn_/jit_fn_/Mul:0' shape=() dtype=float32>
  _aval = ShapedArray(float32[], weak_type=True), 2*width, 3).
If using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.

If I switch the image.resize to use an integer, e.g. upscale (h = h * 2), all works fine. However, I wish to downscale the image (h = h / 2), so I need to use a float.

Is there any way to solve this?

[gnecula]

At some level this is independent of jax2tf: dimensions must be integers. I.e., can you jit the function that you are trying to convert with jax2tf? You probably still want to use an integer and use h // 2 for downscaling. Once you start using division you are approaching the limits of the reasoning that shape polymorphism can do in terms of comparing dimension values. You may run into some errors that are specific to shape polymorphism. If you are prepared to ensure that you always will use this function on images with even height, then you can write 2 * height in lieu of height in the polymorphic shape specification.

[thekevinscott]

Thanks for the response, @gnecula! That resolved my problem.

The context for this question is trying to convert a model written in Jax into Tensorflow, and I've no problem ensuring an even size for the incoming image.

In the original model code, the authors are multiplying by a float and then casting to an int:

https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py#L86

I'd originally tried astype(int), which didn't work. I wonder if there could be someway to enable int() behavior for this use case?

No worries if not. My problem is resolved 👍 Thanks again!

[gnecula]

In this case you cannot cast to an int because the dimension is actually not known at lowering time.

[thekevinscott]

Gotcha. Thanks for the feedback and the help