Best practice of type for a scalar value

[ToshiyukiBandai]

Hi JAX community,

In my code, there are some physical constants. I am wondering what is the best practice of the type for scalar values. I think there are three options:

1. Python float (e.g., x = 2.0),
2. 0D jnp.array (e.g., x = jnp.array(2.0)),
3. 1D jnp.array (e.g., x = jnp.array([2.0])).

It seems to me that the last option is the most computationally effective because it can avoid type promotion when interacting with other 1D array objects. Does anyone have specific recommendations?

[jakevdp]

All of these should be fine. There are several differences:

• outside a JAX-transformed function, python floats will lead to computations being executed on host rather than on device. For example, x = 1.0; y = 2.0; z = x * y is a host-side multiplication, not a multiplication on device.
• In JAX's current implementation, both 1 and 2 will be considered weakly-typed floats, because their dtype is not specified. So if you do x * jnp.bfloat16(1.0), the result will be bfloat16 rather than promoting the bfloat32 to float32 or float64
• The third option is strongly-typed under JAX's current implementation, because for list inputs JAX first converts to numpy array for efficiency. So if you do x * jnp.bfloat16(1.0), the result will be float32 or float64. This may or may not be desirable depending on your use-case. An alternative way to have a strongly-typed scalar is to specify it more explicitly; e.g. x = jnp.float32(2.0)

Does that answer your question?

[ToshiyukiBandai]

Hi @jakevdp

Thank you for answering my question! So, in terms of computational time, type promotion itself is quite negligible, right?

Regarding your first comment, does that mean I should avoid using Python floats when running the program on GPU? I thought that JAX needs to ask CPU when JAX finds Python floats, which would take some additional computational time. Maybe, I am confused with host and device...

[jakevdp]

No, using python floats is fine. If you do, e.g. 2 * jnp.asarray(x), that computation will happen on GPU. If you pass them to a jax function, like jnp.add(2, 2), that computation will happen on GPU. It's only if you do a Python expression that does not involve JAX, like 2 + 2, that it will happen on-host.

[jakevdp]

From a runtime perspective, I don't think this is worth worrying about.

[ToshiyukiBandai]

Thank you! It makes sense. I appreciate your help!