improve vmap_method description

Author: rdyro

torch2jax/api.py

@@ -150,9 +150,12 @@ def torch2jax(
         example_kw: Example keyword arguments. Defaults to None.
         output_shapes: Output shapes or shapes + dtype struct. Defaults to None.
         output_sharding_spec: jax.sharding.PartitionSpec specifying the sharding spec of the output, uses input mesh.
-        vmap_method: batching method, see https://docs.jax.dev/en/latest/ffi.html#batching-with-vmap 
+        vmap_method: batching method, see 
+            [https://docs.jax.dev/en/latest/ffi.html#batching-with-vmap](https://docs.jax.dev/en/latest/ffi.html#batching-with-vmap)
+
             NOTE: only vmap_method="sequntial" is supported non-experimentally
-            NOTE: try "expand_dims", "broadcast_all" if you want experimentally use pytorch-side batching
+
+            NOTE: try "expand_dims", "broadcast_all" if you want to experiment with pytorch-side batching
     Returns:
         Callable: JIT-compatible JAX function.
 

torch2jax/gradients.py

@@ -46,9 +46,12 @@ def torch2jax_with_vjp(
             library PyTorch code may need this fallback. Defaults to True (i.e., do not use fallback).
         output_sharding_spec: (not supported) sharding spec of the output, use shard_map instead for a device-local
             version of this function
-        vmap_method: batching method, see https://docs.jax.dev/en/latest/ffi.html#batching-with-vmap 
+        vmap_method: batching method, see 
+            [https://docs.jax.dev/en/latest/ffi.html#batching-with-vmap](https://docs.jax.dev/en/latest/ffi.html#batching-with-vmap)
+
             NOTE: only vmap_method="sequntial" is supported non-experimentally
-            NOTE: try "expand_dims", "broadcast_all" if you want experimentally use pytorch-side batching
+
+            NOTE: try "expand_dims", "broadcast_all" if you want to experiment with pytorch-side batching
     Returns:
         Callable: JIT-compatible JAX version of the torch function (VJP defined up to depth `depth`).