Add support for callable in torchax.interop.JittableModule.functional_call in the first parameter

torchax/test/test_jittable_module.py

@@ -34,6 +34,25 @@ def test_isinstance_does_not_mix(self):
     assert isinstance(JittableMoreAwesomeModel, EvenMoreAwesomeModel)
     assert not isinstance(JittableMoreAwesomeModel, MyAwesomeModel)
 
+  def test_functional_call_callable(self):
+
+    def outer_function(model, x):
+      return x + 1
+
+    model = MyAwesomeModel()
+    jittable_module = interop.JittableModule(model)
+
+    # Check if the jittable module can be called like a function
+    input_tensor = torch.randn(1, 3, 224, 224)
+    expected_output = input_tensor + 1
+
+    output = jittable_module.functional_call(outer_function,
+                                             jittable_module.params,
+                                             jittable_module.buffers,
+                                             input_tensor)
+
+    assert torch.equal(output, expected_output)
+
 
 if __name__ == '__main__':
   unittest.main()

torchax/torchax/interop.py

@@ -90,16 +90,26 @@ def __class__(self):
   def __call__(self, *args, **kwargs):
     return self.forward(*args, **kwargs)
 
-  def functional_call(self, method_name, params, buffers, *args, **kwargs):
+  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 = getattr(self._model, method_name)(*args, **kwargs)
+      res = method(*args, **kwargs)
     return res
 
   def forward(self, *args, **kwargs):