folx 0.2.2__tar.gz → 0.2.2.post1__tar.gz

{folx-0.2.2 → folx-0.2.2.post1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: folx
-Version: 0.2.2
+Version: 0.2.2.post1
 Summary: Forward Laplacian for JAX
 Home-page: https://github.com/microsoft/folx
 License: MIT

{folx-0.2.2 → folx-0.2.2.post1}/folx/vmap.py

@@ -1,6 +1,5 @@
 import functools
-import logging
-from typing import Any, Callable, Sequence, TypeVar, cast
+from typing import Any, Callable, Sequence, TypeVar
 
 import jax
 import jax.numpy as jnp
@@ -18,6 +17,7 @@ def batched_vmap(
     if isinstance(in_axes, list):
         in_axes = tuple(in_axes)
     
+    @functools.wraps(fn)
     def result(*args, **kwargs):
         wrapped_fn = functools.partial(fn, **kwargs)
 
@@ -41,22 +41,24 @@ def batched_vmap(
         if batch_size <= max_batch_size:
             return jax.vmap(wrapped_fn, in_axes=in_axes, out_axes=out_axes)(*args)
         
-        if batch_size % max_batch_size != 0:
-            logging.warning(f"Batch size {batch_size} is not divisible by max_batch_size {max_batch_size}. Switching to vmap.")
-            return jax.vmap(wrapped_fn, in_axes=in_axes, out_axes=out_axes)(*args)
-        
         # Split into batches.
-        assert batch_size % max_batch_size == 0
         num_batches = batch_size // max_batch_size
+        remainder = batch_size % max_batch_size
         leading_args = [
-            jtu.tree_map(lambda x: jnp.moveaxis(x, ax, 0).reshape(-1, max_batch_size, *x.shape[1:]), arg)
+            jtu.tree_map(lambda x: jnp.moveaxis(x, ax, 0), arg)
             for arg, ax in zip(mapped_args, mapped_axes)
         ]
+        loop_args = jtu.tree_map(lambda x: x[remainder:].reshape(num_batches, max_batch_size, *x.shape[1:]), leading_args)
+
+        vmapped_fn = jax.vmap(lambda x: wrapped_fn(*in_tree.unflatten(merge(x))), in_axes=0, out_axes=0)
         def inner(carry, x):
-            vmapped_fn = jax.vmap(lambda x: wrapped_fn(*in_tree.unflatten(merge(x))), in_axes=0, out_axes=0)
             return carry, vmapped_fn(x)
-        result = jax.lax.scan(inner, None, leading_args, length=num_batches)[1]
+        result = jax.lax.scan(inner, None, loop_args, length=num_batches)[1]
         result = jtu.tree_map(lambda x: x.reshape(-1, *x.shape[2:]), result)
+        if remainder > 0:
+            remainder_args = jtu.tree_map(lambda x: x[:remainder], leading_args)
+            remainder_result = vmapped_fn(remainder_args)
+            result = jtu.tree_map(lambda x, r: jnp.concatenate([r, x], axis=0), result, remainder_result)
 
         out_axes_flat, out_tree = jtu.tree_flatten(out_axes)
         result = tuple(
@@ -64,4 +66,4 @@ def batched_vmap(
             for r, ax in zip(out_tree.flatten_up_to(result), out_axes_flat)
         )
         return out_tree.unflatten(result)
-    return cast(F, result)
+    return result

{folx-0.2.2 → folx-0.2.2.post1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = 'folx'
-version = '0.2.2'
+version = '0.2.2-1'
 description = 'Forward Laplacian for JAX'
 authors = [
     "Nicholas Gao <n.gao@tum.de>",