@jtu.skip_on_devices('tpu')
def testOnesBroadcastingConstantHandler(self):
    self.skipTest('test needs jax3 update')

    def fun(x):
        ones = jnp.ones((3, 4))
        assert (isinstance(ones, onp.ndarray) and (ones.strides == (0, 0)))
        c = x._node.c
        Broadcast = c.Broadcast
        was_called = []
        c.Broadcast = (lambda *args: (was_called.append(True) or Broadcast(*args)))
        out = (x + ones)
        assert was_called, 'Broadcast was not called.'
        return out
    fun = api.jit(fun)
    out_val = fun(jnp.ones(4))
    self.assertAllClose(out_val, onp.full((3, 4), 2.0), check_dtypes=False)
@jtu.skip_on_devices('tpu')
def testOnesBroadcastingConstantHandler(self):
    self.skipTest('test needs jax3 update')

    def fun(x):
        ones = jnp.ones((3, 4))
        assert (isinstance(ones, onp.ndarray) and (ones.strides == (0, 0)))
        c = x._node.c
        Broadcast = c.Broadcast
        was_called = []
        c.Broadcast = (lambda *args: (was_called.append(True) or Broadcast(*args)))
        out = (x + ones)
        assert was_called, 'Broadcast was not called.'
        return out
    fun = api.jit(fun)
    out_val = fun(jnp.ones(4))
    self.assertAllClose(out_val, onp.full((3, 4), 2.0), check_dtypes=False)