def METHOD_NAME(self, stateless):
    if not stateless and backend.backend() == "jax":
        self.skipTest("LossScaleOptimizer must use stateless_apply on jax.")
    inner_optimizer = SGD(learning_rate=0.5)
    optimizer = LossScaleOptimizer(inner_optimizer, initial_scale=400.0)
    vars = [backend.Variable([1.0, 2.0, 3.0, 4.0])]
    optimizer.build(vars)
    opt_vars = optimizer.variables
    grads = [ops.array([np.inf, np.inf, np.inf, np.inf])]
    for _ in range(4):
        if stateless:
            _, opt_vars = optimizer.stateless_apply(opt_vars, grads, vars)
            for ref_v, v in zip(optimizer.variables, opt_vars):
                ref_v.assign(v)
        else:
            optimizer.apply(grads, vars)
    self.assertAllClose(optimizer.scale_loss(1.0), 25.0)