code: def test_nn_module(batch_shape):
      import torch
  
      class Matmul(torch.nn.Module):
          def forward(
              self, x, y
          ) -> Dependent[lambda x, y: Reals[x.shape[0], y.shape[1]]]:
              return x @ y
  
      matmul = make_op(Matmul())
  
      inputs = OrderedDict(
          (k, Bint[s]) for k, s in zip("abc", batch_shape)
      )
      for L, J, K in [(2, 3, 4), (5, 6, 7)]:
          # Check symbolic.
          x = Variable("x", Reals[L, J])
          y = Variable("y", Reals[J, K])
          actual = matmul(x, y)
          assert isinstance(actual, Binary)
          assert actual.output == Reals[L, K]
  
          # Check on ground term.
          x_data = randn(*batch_shape, L, J)
          y_data = randn(*batch_shape, J, K)
          actual_data = matmul(x_data, y_data)
          assert isinstance(actual_data, type(x_data))
          assert_close(actual_data, x_data @ y_data)
  
          # Check on funsor.Tensors.
          x = Tensor(x_data, inputs)
          y = Tensor(y_data, inputs)
          actual = matmul(x, y)
          assert isinstance(actual, Tensor)
          assert_close(actual, Tensor(actual_data, inputs))
  
file_path: funsor/test/test_op_factory.py
imports: from collections import OrderedDict
  import pytest
  from funsor.domains import Array, Bint, Dependent, Product, Reals
  from funsor.op_factory import make_op
  from funsor.tensor import Tensor
  from funsor.terms import Binary, Tuple, Unary, Variable
  from funsor.testing import assert_close, randn, requires_backend
label_desc: Use image node_id 4 for calling a global function with example usage: test_nn_module(batch_shape) and returns: unknown