Make a Pipeable Function Mapped

stamox.core.maps.pipe_vmap(func: Callable[..., ~T] = None, *, in_axes = 0, out_axes = 0, axis_name: Hashable = None, axis_size: int = None, name: str = None) -> Callable[..., ~T]

Creates a functional from a function with vmap.

Parameters:

Name	Type	Description	Default
func	Callable[..., ~T]	The function to be wrapped.	None
in_axes		The number of input axes.	0
out_axes		The number of output axes.	0
axis_name	Hashable	The name of the axis.	None
axis_size	int	The size of the axis.	None
name	str	The name of the functional. If not provided, the name of the	None

Returns:

Type	Description
Callable[..., ~T]	A callable that creates a functional from the given function.

Examples:

>>> from stamox import pipe_vmap
>>> f = lambda x: x + 1
>>> f = pipe_vmap(f)
>>> g = f >> f >> f
>>> g(jnp.array([1, 2, 3]))
Array([4, 5, 6], dtype=int32)

---

stamox.core.maps.pipe_pmap(func: Callable[..., ~T] = None, *, in_axes = 0, out_axes = 0, axis_name: Hashable = None, axis_size: int = None, name: str = None) -> Callable[..., ~T]

Creates a functional object from a given function.

Parameters:

Name	Type	Description	Default
func	Callable[P, T]	The function to be wrapped.	None
in_axes	int	The number of input axes for the function.	0
out_axes	int	The number of output axes for the function.	0
axis_name	Hashable	The name of the axis.	None
axis_size	int	The size of the axis.	None
name	str	The name of the functional object.	None

Returns:

Type	Description
Callable[P, T]	A callable object that wraps the given function.

Examples:

>>> from stamox import pipe_pmap
>>> f = lambda x: x + 1
>>> f = pipe_pmap(f)
>>> g = f >> f >> f
>>> g(jnp.array([1, 2, 3]))
Array([4, 5, 6], dtype=int32)

---

stamox.core.maps.partial_pipe_vmap(func: Callable[..., ~T] = None, *, name: str = None) -> Callable[..., ~T]

Partially apply a function to a vmap.

Parameters:

Name	Type	Description	Default
func	Callable[P, T]	The function to partially apply.	None
name	str	The name of the function. Defaults to None.	None

Returns:

Type	Description
Callable[P, T]	A partially applied function.

Examples:

>>> from stamox import partial_pipe_vmap
>>> f = lambda x, y: x + y
>>> f = partial_pipe_vmap(f)
>>> g = f(y=1) >> f(y=2) >> f(y=3)
>>> g(jnp.array([1, 2, 3]))
Array([7, 8, 9], dtype=int32)

---

stamox.core.maps.partial_pipe_pmap(func: Callable[..., ~T] = None, *, name: str = None) -> Callable[..., ~T]

Partially apply a function to a pipe.

Parameters:

Name	Type	Description	Default
func	Callable[P, T]	The function to partially apply.	None
name	str	The name of the function. Defaults to None.	None

Returns:

Type	Description
Callable[P, T]	A partially applied function.

Examples:

>>> from stamox import partial_pipe_pmap
>>> f = lambda x, y: x + y
>>> f = partial_pipe_pmap(f)
>>> g = f(y=1) >> f(y=2) >> f(y=3)
>>> g(jnp.array([1, 2, 3]))
Array([7, 8, 9], dtype=int32)