Would nn.BatchNorm2d work as expected?

[lime-j]

Dear authors,

Thanks for your wonderful work! I wonder if nn.BatchNorm2d works. During training, the running mean and running variance are updated according to their own rule instead of being updated with the optimizer.

[Artur-Galstyan]

Hey, I wanted to check this out tonight, I'll let you know how it goes :)

[lime-j]

wow, thanks for your attention :)

[Artur-Galstyan]

@lime-j Ok so I tested a few things and I don't know exactly what you mean by "if nn.BatchNorm2d works". AFAICT, it seems to work just fine. See this script, which "trains" a model that includes batch norm using Equinox and Optax and compares the gradients with the JAX ones.

from typing import Any

import equinox as eqx
import jax
import jax.numpy as jnp
import numpy as np
import optax
import torch
import torchax
from jaxtyping import PyTree
from torchax.interop import JittableModule, call_torch, jax_view

torchax.enable_globally()


class MinimalTorchChain(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.norm = torch.nn.BatchNorm2d(num_features=1)
        self.conv1 = torch.nn.Conv2d(
            in_channels=1, out_channels=2, kernel_size=3, padding=1
        )
        self.pool = torch.nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
        self.conv2 = torch.nn.Conv2d(
            in_channels=2, out_channels=1, kernel_size=1, padding=1
        )

    def forward(self, x):
        x = self.norm(x)
        x = self.conv1(x)
        x = self.pool(x)
        x = self.conv2(x)
        return x


class HybridChainModel(eqx.Module):
    params: PyTree
    _call: Any

    def __init__(self):
        torch_model = MinimalTorchChain()
        torch_model.to("jax")
        torch_jitted = JittableModule(torch_model)
        self.params = jax_view(torch_jitted.params)

        def _call_model(params, x):
            return torch_jitted.functional_call(
                "forward", params, torch_jitted.buffers, x
            )

        self._call = _call_model

    def __call__(self, x):
        return call_torch(self._call, self.params, x)


model = HybridChainModel()
key = jax.random.key(0)
dummy_x = jax.random.normal(key, (2, 1, 4, 4))

original_params = {k: np.array(v) for k, v in model.params.items()}

optimizer = optax.adam(learning_rate=0.1)
opt_state = optimizer.init(eqx.filter(model, eqx.is_array))


def loss_fn(m, x):
    pred = m(x)
    return jnp.mean(pred)


print("batch norm before", model.params["norm.weight"])
loss_val, grads = eqx.filter_value_and_grad(loss_fn)(model, dummy_x)
print(grads.params)
updates, opt_state = optimizer.update(grads, opt_state, model)
model = eqx.apply_updates(model, updates)
print("batch norm after", model.params["norm.weight"])

torch_model_ref = MinimalTorchChain()
torch_model_ref.to("cpu")

with torch.no_grad():
    torch_model_ref.norm.weight.copy_(
        torch.tensor(original_params["norm.weight"]).reshape(-1)
    )
    torch_model_ref.norm.bias.copy_(
        torch.tensor(original_params["norm.bias"]).reshape(-1)
    )
    torch_model_ref.conv1.weight.copy_(torch.tensor(original_params["conv1.weight"]))
    torch_model_ref.conv1.bias.copy_(torch.tensor(original_params["conv1.bias"]))
    torch_model_ref.conv2.weight.copy_(torch.tensor(original_params["conv2.weight"]))
    torch_model_ref.conv2.bias.copy_(torch.tensor(original_params["conv2.bias"]))

for param in torch_model_ref.parameters():
    param.requires_grad = True

torch_x = torch.tensor(np.array(dummy_x), requires_grad=False)
torch_output = torch_model_ref(torch_x)
torch_loss = torch_output.mean()
torch_loss.backward()

print("\n=== GRADIENT COMPARISON ===")
for name in [
    "norm.weight",
    "norm.bias",
    "conv1.weight",
    "conv1.bias",
    "conv2.weight",
    "conv2.bias",
]:
    jax_grad = np.array(grads.params[name])

    layer_name = name.split(".")[0]
    param_type = name.split(".")[1]

    if param_type == "weight":
        torch_grad = np.array(getattr(torch_model_ref, layer_name).weight.grad)
    else:
        torch_grad = np.array(getattr(torch_model_ref, layer_name).bias.grad)

    diff = np.abs(jax_grad - torch_grad).max()
    print(f"{name}: max diff = {diff:.2e}")

When comparing the gradients, we can see this:

(torchax-test) arturgalstyan@talon:~/Workspace/torchax-test$ python3 test8.py
WARNING:absl:Tensorflow library not found, tensorflow.io.gfile operations will use native shim calls. GCS paths (i.e. 'gs://...') cannot be accessed.
WARNING:root:Duplicate op registration for aten.__and__
batch norm before [1.]
{'conv1.bias': Array([-0.04179712,  0.05805584], dtype=float32), 'conv1.weight': Array([[[[-0.02758748, -0.00026355,  0.0024729 ],
         [ 0.04498183, -0.01467221,  0.01483689],
         [ 0.01460248, -0.02848585,  0.00789216]]],


       [[[ 0.01670449,  0.02933157, -0.01356538],
         [ 0.02416716, -0.03358139, -0.04470449],
         [ 0.01890323,  0.00251995,  0.04067027]]]], dtype=float32), 'conv2.bias': Array([1.], dtype=float32), 'conv2.weight': Array([[[[0.4383718 ]],

        [[0.34648165]]]], dtype=float32), 'norm.bias': Array([0.02774466], dtype=float32), 'norm.weight': Array([0.00820458], dtype=float32)}
batch norm after [0.9000008]

=== GRADIENT COMPARISON ===
/home/arturgalstyan/Workspace/torchax-test/test8.py:117: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. To learn more, see the migration guide https://numpy.org/devdocs/numpy_2_0_migration_guide.html#adapting-to-changes-in-the-copy-keyword
  torch_grad = np.array(getattr(torch_model_ref, layer_name).weight.grad)
norm.weight: max diff = 1.86e-09
/home/arturgalstyan/Workspace/torchax-test/test8.py:119: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. To learn more, see the migration guide https://numpy.org/devdocs/numpy_2_0_migration_guide.html#adapting-to-changes-in-the-copy-keyword
  torch_grad = np.array(getattr(torch_model_ref, layer_name).bias.grad)
norm.bias: max diff = 0.00e+00
conv1.weight: max diff = 1.86e-09
conv1.bias: max diff = 3.73e-09
conv2.weight: max diff = 1.19e-07
conv2.bias: max diff = 1.19e-07
(torchax-test) arturgalstyan@talon:~/Workspace/torchax-test$

So I would say, yes, batch norm works. Is this what you meant?

[lime-j]

during training, the nn.BatchNorm2d have its own update rule for some of its parameters. It have a running mean, running var, and keeps a num_batches_tracked.

[lime-j]

during training, the nn.BatchNorm2d have its own update rule for some of its parameters. It have a running mean, running var, and keeps a num_batches_tracked.

these torch param are not updated by optimizer. but once you seperate the model and the param and convert them into the jax view (as HybridChainModel in your code), the jitted model function doesn't seems to be able for updating the param properly.

[Artur-Galstyan]

Ahhh, now I see what you mean. And yes, you are onto something. See this code:

from typing import Any

import equinox as eqx
import jax
import torch
import torchax
from jaxtyping import PyTree
from torchax.interop import JittableModule, call_torch, jax_view

torchax.enable_globally()


class MinimalTorchChain(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.norm = torch.nn.BatchNorm2d(num_features=1)
        self.conv1 = torch.nn.Conv2d(
            in_channels=1, out_channels=2, kernel_size=3, padding=1
        )
        self.pool = torch.nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
        self.conv2 = torch.nn.Conv2d(
            in_channels=2, out_channels=1, kernel_size=1, padding=1
        )

    def forward(self, x):
        x = self.norm(x)
        x = self.conv1(x)
        x = self.pool(x)
        x = self.conv2(x)
        return x


class HybridChainModel(eqx.Module):
    params: PyTree
    _torch_jitted: Any
    _call: Any

    def __init__(self):
        torch_model = MinimalTorchChain()
        torch_model.train()
        print("Torch model training mode:", torch_model.training)
        torch_model.to("jax")
        torch_jitted = JittableModule(torch_model)
        self._torch_jitted = torch_jitted
        self.params = jax_view(torch_jitted.params)

        def _call_model(params, x):
            returned = self._torch_jitted.functional_call(
                "forward", params, self._torch_jitted.buffers, x
            )
            print(f"{returned=}")
            print(f"{self._torch_jitted.buffers=}")
            return returned

        self._call = _call_model

    def __call__(self, x):
        return call_torch(self._call, self.params, x)


model = HybridChainModel()
key = jax.random.key(0)
dummy_x = jax.random.normal(key, (2, 1, 4, 4))

print("=== BEFORE FORWARD ===")
print("torch_jitted.buffers:", model._torch_jitted.buffers)

output = model(dummy_x)

print("\n=== AFTER FORWARD ===")
print("torch_jitted.buffers:", model._torch_jitted.buffers)

It gives this output:

(torchax-test) arturgalstyan@talon:~/Workspace/torchax-test$ python3 test10.py
WARNING:absl:Tensorflow library not found, tensorflow.io.gfile operations will use native shim calls. GCS paths (i.e. 'gs://...') cannot be accessed.
WARNING:root:Duplicate op registration for aten.__and__
Torch model training mode: True
=== BEFORE FORWARD ===
torch_jitted.buffers: {'norm.num_batches_tracked': Tensor(<class 'jaxlib._jax.ArrayImpl'> 0), 'norm.running_mean': Tensor(<class 'jaxlib._jax.ArrayImpl'> [0.]), 'norm.running_var': Tensor(<class 'jaxlib._jax.ArrayImpl'> [1.])}
returned=Tensor(<class 'jaxlib._jax.ArrayImpl'> [[[[0.6692908  0.6692908  0.6692908  0.6692908  0.6692908  0.6692908 ]
   [0.6692908  0.9719087  0.9719087  0.96411073 0.86056733 0.6692908 ]
   [0.6692908  0.9719087  0.9719087  0.96411073 0.86056733 0.6692908 ]
   [0.6692908  0.9719087  0.9719087  0.96411073 0.93189085 0.6692908 ]
   [0.6692908  0.65900207 0.66188765 0.93189085 0.93189085 0.6692908 ]
   [0.6692908  0.6692908  0.6692908  0.6692908  0.6692908  0.6692908 ]]]


 [[[0.6692908  0.6692908  0.6692908  0.6692908  0.6692908  0.6692908 ]
   [0.6692908  0.76938176 0.76938176 0.73408157 0.6854578  0.6692908 ]
   [0.6692908  0.7818705  0.82939667 0.82939667 0.7980425  0.6692908 ]
   [0.6692908  0.7077926  0.82939667 0.82939667 0.8162132  0.6692908 ]
   [0.6692908  0.7077926  0.82939667 0.82939667 0.8162132  0.6692908 ]
   [0.6692908  0.6692908  0.6692908  0.6692908  0.6692908  0.6692908 ]]]])
self._torch_jitted.buffers={'norm.num_batches_tracked': Tensor(<class 'jaxlib._jax.ArrayImpl'> 1), 'norm.running_mean': Tensor(<class 'jaxlib._jax.ArrayImpl'> [0.]), 'norm.running_var': Tensor(<class 'jaxlib._jax.ArrayImpl'> [1.])}

=== AFTER FORWARD ===
torch_jitted.buffers: {'norm.num_batches_tracked': Tensor(<class 'jaxlib._jax.ArrayImpl'> 1), 'norm.running_mean': Tensor(<class 'jaxlib._jax.ArrayImpl'> [0.]), 'norm.running_var': Tensor(<class 'jaxlib._jax.ArrayImpl'> [1.])}
(torchax-test) arturgalstyan@talon:~/Workspace/torchax-test$

As you can see, the norm.num_batches_tracked increments, but the running_var and running_mean don't! I will investigate this to see if this can be fixed. The issue is that the functional call only returns the tensors of actually trainable params (not buffers). Ideally, it should probably return a state, tensor tuple instead. I will see what I can do here.

[Artur-Galstyan]

Ok, so this is how far I got:

@op(torch.ops.aten._native_batch_norm_legit.default)
def _aten__native_batch_norm_legit(
    input, weight, bias, running_mean, running_var, training, momentum, eps
):
    ...

    return x_hat, saved_mean, saved_rstd # <<< does not return the running mean and var!

This function doesn't return the running mean and var, but if you tried, you would get a segfault, likely because PyTorch does not return it either. PyTorch mutates those in place but we can't do that in JAX. We'd need a way to return those, tell PyTorch to ignore it and in the functional_call return this somehow:

class JittableModule(torch.nn.Module):
  ...

  def functional_call(self, method_or_name, params, buffers, *args, **kwargs):
      kwargs = kwargs or {}
      params_copy = copy.copy(params)
      params_copy.update(buffers)
      # reinflate the state dict so there are not any missing keys
      for k, v in self._extra_dumped_weights.items():
        for new_key in v:
          params_copy[new_key] = params_copy[k]

      if isinstance(method_or_name, str):
        method = getattr(self._model, method_or_name)
      else:
        if not callable(method_or_name):
          raise TypeError(
              f"method_or_name should be a callable or a string, got {type(method_or_name)}"
          )
        method = method_or_name
        args = (self._model,) + args
      with torch_stateless._reparametrize_module(self._model, params_copy):
          res = method(*args, **kwargs)

          updated_buffers = {}
          for name, buffer in self._model.named_buffers():
              updated_buffers[name] = buffer
      return res, updated_buffers # return both as a <tensor, state> tuple. 
...

@qihqi What do you think?

[Artur-Galstyan]

@qihqi @lime-j I think I got it working (albeit with lots of AI help):

@op(torch.ops.aten._native_batch_norm_legit.default, is_jax_function=False)
def _aten__native_batch_norm_legit(
    input, weight, bias, running_mean, running_var, training, momentum, eps
):
    env = input._env
    input_jax = input.jax()
    weight_jax = weight.jax() if weight is not None else None
    bias_jax = bias.jax() if bias is not None else None
    rm_jax = running_mean.jax()
    rv_jax = running_var.jax()

    reduction_dims = [0] + list(range(2, input_jax.ndim))
    reshape_dims = [1, -1] + [1] * (input_jax.ndim - 2)
    if training:
        mean = jnp.mean(input_jax, axis=reduction_dims, keepdims=True)
        saved_mean = jnp.squeeze(mean, reduction_dims)
        var = jnp.var(input_jax, axis=reduction_dims)
        rstd = jax.lax.rsqrt(var.reshape(reshape_dims) + eps)

        new_running_mean = (1 - momentum) * rm_jax + momentum * saved_mean
        new_running_var = (1 - momentum) * rv_jax + momentum * var

        running_mean._elem = new_running_mean
        running_var._elem = new_running_var

        saved_rstd = jnp.squeeze(rstd, reduction_dims)

        x_hat = (input_jax - mean) * rstd
    else:
        rstd = jax.lax.rsqrt(rv_jax.reshape(reshape_dims) + eps)
        saved_mean = jnp.array(
            [], dtype=input_jax.dtype
        )
        saved_rstd = jnp.array([], dtype=input_jax.dtype)
        x_hat = (input_jax - rm_jax.reshape(reshape_dims)) * rstd

    if weight_jax is not None:
        x_hat *= weight_jax.reshape(reshape_dims)
    if bias_jax is not None:
        x_hat += bias_jax.reshape(reshape_dims)

    return (
        torchax.tensor.Tensor(x_hat, env),
        torchax.tensor.Tensor(saved_mean, env),
        torchax.tensor.Tensor(saved_rstd, env),
    )
...

@op(torch.ops.aten.native_batch_norm, is_jax_function=False)
def _aten_native_batch_norm(
    input,
    weight,
    bias,
    running_mean,
    running_var,
    training=False,
    momentum=0.1,
    eps=1e-5,
):
    if running_mean is None:
        rm_jax = jnp.zeros(input.shape[1], dtype=input.jax().dtype)
        running_mean = torchax.tensor.Tensor(rm_jax, input._env)
    if running_var is None:
        rv_jax = jnp.ones(input.shape[1], dtype=input.jax().dtype)
        running_var = torchax.tensor.Tensor(rv_jax, input._env)

    if training:
        return _aten__native_batch_norm_legit(
            input,
            weight,
            bias,
            running_mean,
            running_var,
            training,
            momentum,
            eps,
        )
    else:
        return _aten__native_batch_norm_legit_no_training(
            input, weight, bias, running_mean, running_var, momentum, eps
        )

...

@op(torch.ops.aten._native_batch_norm_legit_no_training, is_jax_function=False)
def _aten__native_batch_norm_legit_no_training(
    input, weight, bias, running_mean, running_var, momentum, eps
):
    return _aten__native_batch_norm_legit(
        input, weight, bias, running_mean, running_var, False, momentum, eps
    )

Basically, just added is_jax_function=False but I'm really not sure if this is the correct way. However, testing it with my other testing file:

from typing import Any

import equinox as eqx
import jax
import jax.numpy as jnp
import numpy as np
import optax
import torch
import torchax
from jaxtyping import PyTree
from torchax.interop import JittableModule, call_torch, jax_view

torchax.enable_globally()


class MinimalTorchChain(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.norm = torch.nn.BatchNorm2d(num_features=1)
        self.conv1 = torch.nn.Conv2d(
            in_channels=1, out_channels=2, kernel_size=3, padding=1
        )
        self.pool = torch.nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
        self.conv2 = torch.nn.Conv2d(
            in_channels=2, out_channels=1, kernel_size=1, padding=1
        )

    def forward(self, x):
        x = self.norm(x)
        x = self.conv1(x)
        x = self.pool(x)
        x = self.conv2(x)
        return x


class HybridChainModel(eqx.Module):
    params: PyTree
    buffers: Any
    _call: Any

    def __init__(self):
        torch_model = MinimalTorchChain()
        torch_model.to("jax")
        torch_jitted = JittableModule(torch_model)
        self.buffers = torch_jitted.buffers
        self.params = jax_view(torch_jitted.params)

        def _call_model(params, x):
            return torch_jitted.functional_call("forward", params, self.buffers, x)

        self._call = _call_model

    def __call__(self, x):
        return call_torch(self._call, self.params, x)


model = HybridChainModel()
key = jax.random.key(0)
dummy_x = jax.random.normal(key, (2, 1, 4, 4))

original_params = {k: np.array(v) for k, v in model.params.items()}

optimizer = optax.adam(learning_rate=0.1)
opt_state = optimizer.init(eqx.filter(model, eqx.is_array))


def loss_fn(m, x):
    pred = m(x)
    return jnp.mean(pred)


print("before", model.buffers)

print("batch norm before", model.params["norm.weight"])
loss_val, grads = eqx.filter_value_and_grad(loss_fn)(model, dummy_x)
print(grads.params)
updates, opt_state = optimizer.update(grads, opt_state, model)
model = eqx.apply_updates(model, updates)
print("batch norm after", model.params["norm.weight"])

torch_model_ref = MinimalTorchChain()
torch_model_ref.to("cpu")

with torch.no_grad():
    torch_model_ref.norm.weight.copy_(
        torch.tensor(original_params["norm.weight"]).reshape(-1)
    )
    torch_model_ref.norm.bias.copy_(
        torch.tensor(original_params["norm.bias"]).reshape(-1)
    )
    torch_model_ref.conv1.weight.copy_(torch.tensor(original_params["conv1.weight"]))
    torch_model_ref.conv1.bias.copy_(torch.tensor(original_params["conv1.bias"]))
    torch_model_ref.conv2.weight.copy_(torch.tensor(original_params["conv2.weight"]))
    torch_model_ref.conv2.bias.copy_(torch.tensor(original_params["conv2.bias"]))

for param in torch_model_ref.parameters():
    param.requires_grad = True

torch_x = torch.tensor(np.array(dummy_x), requires_grad=False)

print(
    "before (torch):",
    torch_model_ref.norm.running_mean,
    torch_model_ref.norm.running_var,
    torch_model_ref.norm.num_batches_tracked,
)

torch_output = torch_model_ref(torch_x)
torch_loss = torch_output.mean()
torch_loss.backward()

print(
    "after (torch):",
    torch_model_ref.norm.running_mean,
    torch_model_ref.norm.running_var,
    torch_model_ref.norm.num_batches_tracked,
)
print("\n=== GRADIENT COMPARISON ===")
for name in [
    "norm.weight",
    "norm.bias",
    "conv1.weight",
    "conv1.bias",
    "conv2.weight",
    "conv2.bias",
]:
    jax_grad = np.array(grads.params[name])

    layer_name = name.split(".")[0]
    param_type = name.split(".")[1]

    if param_type == "weight":
        torch_grad = np.array(getattr(torch_model_ref, layer_name).weight.grad)
    else:
        torch_grad = np.array(getattr(torch_model_ref, layer_name).bias.grad)

    diff = np.abs(jax_grad - torch_grad).max()
    print(f"{name}: max diff = {diff:.2e}")


print("after", model.buffers)

Gives this output:

(torchax-test) arturgalstyan@talon:~/Workspace/torchax-test$ python3 test8.py
WARNING:absl:Tensorflow library not found, tensorflow.io.gfile operations will use native shim calls. GCS paths (i.e. 'gs://...') cannot be accessed.
WARNING:root:Duplicate op registration for aten.__and__
before {'norm.num_batches_tracked': Tensor(<class 'jaxlib._jax.ArrayImpl'> 0), 'norm.running_mean': Tensor(<class 'jaxlib._jax.ArrayImpl'> [0.]), 'norm.running_var': Tensor(<class 'jaxlib._jax.ArrayImpl'> [1.])}
batch norm before [1.]
{'conv1.bias': Array([0.07643214, 0.07761794], dtype=float32), 'conv1.weight': Array([[[[-0.04452283,  0.0729696 , -0.04085746],
         [ 0.03836604,  0.0073434 ,  0.03084425],
         [-0.05455584, -0.02048168, -0.01581647]]],


       [[[ 0.01322007, -0.03351644,  0.06861556],
         [-0.02617224,  0.02487036, -0.00499191],
         [ 0.02036906, -0.04721183, -0.0107778 ]]]], dtype=float32), 'conv2.bias': Array([1.], dtype=float32), 'conv2.weight': Array([[[[0.42913574]],

        [[0.42914578]]]], dtype=float32), 'norm.bias': Array([-0.05925386], dtype=float32), 'norm.weight': Array([0.10704657], dtype=float32)}
batch norm after [0.9000007]
before (torch): tensor([0.]) tensor([1.]) tensor(0)
after (torch): tensor([0.0174]) tensor([1.0267]) tensor(1)

=== GRADIENT COMPARISON ===
/home/arturgalstyan/Workspace/torchax-test/test8.py:133: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. To learn more, see the migration guide https://numpy.org/devdocs/numpy_2_0_migration_guide.html#adapting-to-changes-in-the-copy-keyword
  torch_grad = np.array(getattr(torch_model_ref, layer_name).weight.grad)
norm.weight: max diff = 7.45e-09
/home/arturgalstyan/Workspace/torchax-test/test8.py:135: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. To learn more, see the migration guide https://numpy.org/devdocs/numpy_2_0_migration_guide.html#adapting-to-changes-in-the-copy-keyword
  torch_grad = np.array(getattr(torch_model_ref, layer_name).bias.grad)
norm.bias: max diff = 0.00e+00
conv1.weight: max diff = 7.45e-09
conv1.bias: max diff = 1.49e-08
conv2.weight: max diff = 2.98e-08
conv2.bias: max diff = 1.19e-07
after {'norm.num_batches_tracked': Tensor(<class 'jaxlib._jax.ArrayImpl'> 1), 'norm.running_mean': Tensor(<class 'jaxlib._jax.ArrayImpl'> [0.01737513]), 'norm.running_var': Tensor(<class 'jaxlib._jax.ArrayImpl'> [1.0227048])}
(torchax-test) arturgalstyan@talon:~/Workspace/torchax-test$

So, now it seems to work? But please double check this, I'm not 100% sure, esp. with the gradient bit.

[qihqi]

hi @Artur-Galstyan

Thanks this is great!

Reading it more it seems that the issue is that aten. native_batch_norm is to do in-place modification to the input tensors, but Jax array are immutable?

If so you make the function to take in Tensor subclass instead of Jax array; then swap out the inner jax array with modified jax array; keeping the python reference of then Tensor wrapper the same; thus creating the illusion of in-place modification?

Looking at the above examples, I don't see where you've written the modification back into running_mean and runnign_var; i.e. something like

running_mean._elem = <updated jax array?

Also not sure if the JittableModule need to return modified buffer for jax.jit to capture it

[Artur-Galstyan]

Oh but I did it here:

running_mean._elem = new_running_mean
running_var._elem = new_running_var

[Artur-Galstyan]

@qihqi You're right. We still have one more issue with JIT, namely that it doesn't update the buffers if you JIT the model (it works in eager mode, by setting the _elem directly, but fails in JIT mode):

@op(torch.ops.aten.native_batch_norm, is_jax_function=False, needs_env=True)
def _aten_native_batch_norm(
    input,
    weight,
    bias,
    running_mean,
    running_var,
    training=False,
    momentum=0.1,
    eps=1e-5,
    env=None
):
    if running_mean is None:
        rm_jax = jnp.zeros(input.shape[1], dtype=input.jax().dtype)
        running_mean = torchax.tensor.Tensor(rm_jax, input._env)
    if running_var is None:
        rv_jax = jnp.ones(input.shape[1], dtype=input.jax().dtype)
        running_var = torchax.tensor.Tensor(rv_jax, input._env)

    if training:
        return _aten__native_batch_norm_legit(
            input,
            weight,
            bias,
            running_mean,
            running_var,
            training,
            momentum,
            eps,
            env
        )
    else:
        return _aten__native_batch_norm_legit_no_training(
            input, weight, bias, running_mean, running_var, momentum, eps
        )
...
@op(torch.ops.aten._native_batch_norm_legit.default, is_jax_function=False)
def _aten__native_batch_norm_legit(
    input, weight, bias, running_mean, running_var, training, momentum, eps, env=None
):
    input_jax = input.jax()
    weight_jax = weight.jax() if weight is not None else None
    bias_jax = bias.jax() if bias is not None else None
    rm_jax = running_mean.jax()
    rv_jax = running_var.jax()

    reduction_dims = [0] + list(range(2, input_jax.ndim))
    reshape_dims = [1, -1] + [1] * (input_jax.ndim - 2)
    if training:
        mean = jnp.mean(input_jax, axis=reduction_dims, keepdims=True)
        saved_mean = jnp.squeeze(mean, reduction_dims)
        var = jnp.var(input_jax, axis=reduction_dims)
        rstd = jax.lax.rsqrt(var.reshape(reshape_dims) + eps)

        new_running_mean = (1 - momentum) * rm_jax + momentum * saved_mean
        new_running_var = (1 - momentum) * rv_jax + momentum * var

        running_mean._elem = new_running_mean
        running_var._elem = new_running_var

        saved_rstd = jnp.squeeze(rstd, reduction_dims)

        x_hat = (input_jax - mean) * rstd
    else:
        rstd = jax.lax.rsqrt(rv_jax.reshape(reshape_dims) + eps)
        saved_mean = jnp.array(
            [], dtype=input_jax.dtype
        )
        saved_rstd = jnp.array([], dtype=input_jax.dtype)
        x_hat = (input_jax - rm_jax.reshape(reshape_dims)) * rstd

    if weight_jax is not None:
        x_hat *= weight_jax.reshape(reshape_dims)
    if bias_jax is not None:
        x_hat += bias_jax.reshape(reshape_dims)

    return (
        torchax.tensor.Tensor(x_hat, env),
        torchax.tensor.Tensor(saved_mean, env),
        torchax.tensor.Tensor(saved_rstd, env),
    )

This is my test code:

from typing import Any

import equinox as eqx
import jax
import torch
import torchax
from jaxtyping import PyTree
from torchax.interop import JittableModule, call_torch, jax_view

torchax.enable_globally()


class MinimalTorchChain(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.norm = torch.nn.BatchNorm2d(num_features=1)
        self.conv1 = torch.nn.Conv2d(
            in_channels=1, out_channels=2, kernel_size=3, padding=1
        )
        self.pool = torch.nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
        self.conv2 = torch.nn.Conv2d(
            in_channels=2, out_channels=1, kernel_size=1, padding=1
        )

    def forward(self, x):
        x = self.norm(x)
        x = self.conv1(x)
        x = self.pool(x)
        x = self.conv2(x)
        return x


class HybridChainModel(eqx.Module):
    params: PyTree
    _torch_jitted: Any
    _call: Any

    def __init__(self):
        torch_model = MinimalTorchChain()
        torch_model.train()
        print("Torch model training mode:", torch_model.training)
        torch_model.to("jax")
        torch_jitted = JittableModule(torch_model)
        self._torch_jitted = torch_jitted
        self.params = jax_view(torch_jitted.params)

        def _call_model(params, x):
            tensors = self._torch_jitted.functional_call(
                "forward", params, self._torch_jitted.buffers, x
            )
            return tensors

        self._call = _call_model

    def __call__(self, x):
        return call_torch(self._call, self.params, x)


model = HybridChainModel()
key = jax.random.key(0)
dummy_x = jax.random.normal(key, (2, 1, 4, 4))

# print("=== BEFORE FORWARD ===")
# print("torch_jitted.buffers:", model._torch_jitted.buffers)

output = model(dummy_x)
output = model(dummy_x)

# print("\n=== AFTER FORWARD ===")
# print("torch_jitted.buffers:", model._torch_jitted.buffers)


jitted_model = eqx.filter_jit(model)
output = jitted_model(dummy_x)
output = jitted_model(dummy_x)

(I looked through the code and found needs_env=True, so we don't have to extract it from the input anymore)

I think this is currently at the limit of what I can do here and I don't really know the right solution here. The only thing I can think of is to allow state (e.g. running_mean and running_var) as additional input (a kwarg perhaps?) and also return it. This is the functional, JAX way. But I think this is a bigger change than what I suggested and needs probably some more careful planning (and way more PyTorch knowledge, which I'm lacking xd)

We'd probably need to extend/decouple the batch norm ops a bit from the PyTorch one and then there might be a chance? @qihqi WDYT?

[qihqi]

It might be sufficient if jittable_module returns updated buffers; the UX might be wierd. Let me try something over the weekend.

[reachtarunhere]

Hi is there an update on this? It is critical for training a lot of neural networks