How do I print only the shape with jax.debug.print

[Logon27]

Hello, I am using jax.debug.print in order to print the shape of an array. However when I actually call result.shape it produces...
(array(128, dtype=int32), array(1024, dtype=int32))
Which is a tuple of arrays with the dtype included. And I do not consider dtypes to be shape information. How would I go about getting it to only print...
(128, 1024)
Any help is appreciated. I have tried a lot of different things, but to no avail.

[ASEM000]

You can use the following code:

import jax

import jax.tree_util as jtu


@jtu.register_pytree_node_class
class StaticShape:
    def __init__(self, value: tuple[int, ...]):
        self.value = value

    def tree_flatten(self):
        return (), self.value

    @classmethod
    def tree_unflatten(cls, aux_data, _):
        self = object.__new__(cls)
        self.value = aux_data
        return self

    def __repr__(self):
        return f"{self.value!r}"

    def __hash__(self):
        return hash(self.value)

    def __eq__(self, other):
        return (self.value == other.value) if isinstance(other, StaticShape) else False


@jax.jit
def f(x):
    jax.debug.print("shape={x}", x=StaticShape(x.shape))
    return x


f(jax.numpy.ones([4, 4]))
# shape=(4, 4)
# Array([[1., 1., 1., 1.],
#        [1., 1., 1., 1.],
#        [1., 1., 1., 1.],
#        [1., 1., 1., 1.]], dtype=float32)

I suspect this issue is due to jax.jit changing your tuple of ints to a tuple of arrays.

[Logon27]

This should work fine for my use case. A bit unfortunate that I need a custom class to achieve proper shape printing, but it is what it is. Thank you very much for the help.

[ASEM000]

jax team will introduce jtu.register_static to simplify this process by next release.

[jakevdp]

jax.debug.print is designed for printing dynamic values. Array shapes are always static values, so the better tool would be a standard Python print:

# print dynamic array contents with debug print:
jax.debug.print("{}", results)

# print static array attributes with Python print:
print(results.shape)
print(results.dtype)

[jakevdp]

That's the difference between trace-time and runtime: it's pretty well covered here I think? https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html

[ASEM000]

it's pretty well covered here I think?

Yes,

in the context of jax.debug.print docs

My suggestion is that, if to suggest print inside the docs, maybe point out this behavior also within the jax.debug.print docs as a direct comparison between the approaches.

[Logon27]

@ASEM000 Ah I see what you mean about the difference in the number of prints. I may still want your solution then. I guess I will have to test and see if I think a single print is sufficient in the context of my own program. Because the behavior might be odd if I mix and match both regular prints and jax.debug.print. Because it would display the regular print on the first iteration, but not on subsequent iterations while jax.debug.print would. Which might be confusing. A pretty niche issue. I think both your answers are correct. It just depends on which behavior is sufficient for what the individual developer is trying to do.

[jakevdp]

Thanks - makes sense.

Note that if you need the static quantity to print at every loop iteration, you could always do something like this:

jax.debug.print(f"{arr.shape}")

Since it's a static quantity, you can put it in the static string argument that you pass to debug.print.

Defining and registering a pytree just to print a shape seems like a little too heavy of a solution here, if I'm being honest.

[ASEM000]

Defining and registering a pytree just to print a shape seems like a little too heavy of a solution here, if I'm being honest.

I think you are absolutely right, I'm guilty of overdosing on the use of pytrees :D