# Path: jaxopt/_src/tree_util.py
# def broadcast_pytrees(*trees):
# def tree_scalar_mul(scalar, tree_x):
# def tree_add_scalar_mul(tree_x, scalar, tree_y):
# def _vdot_safe(a, b):
# def tree_vdot(tree_x, tree_y):
# def tree_dot(tree_x, tree_y):
# def tree_sum(tree_x):
# def tree_l2_norm(tree_x, squared=False):
# def tree_zeros_like(tree_x):
# def tree_ones_like(tree_x):
# def tree_average(trees, weights):
# def tree_gram(a):
# def tree_inf_norm(tree_x):
# def tree_where(cond, a, b):
# def tree_negative(tree):
# def tree_reciproqual(tree):
# def tree_mean(tree):


from typing import Any
from typing import Optional
from typing import Tuple
from jaxopt._src import tree_util
import jax
import jax.numpy as jnp
  """
  del hyperparams, scaling
  return x


def prox_lasso(x: Any,
               l1reg: Optional[Any] = None,
               scaling: float = 1.0) -> Any:
  r"""Proximal operator for the l1 norm, i.e., soft-thresholding operator.

  .. math::

    \underset{y}{\text{argmin}} ~ \frac{1}{2} ||x - y||_2^2
    + \text{scaling} \cdot \text{l1reg} \cdot ||y||_1

  When ``l1reg`` is a pytree, the weights are applied coordinate-wise.

  Args:
    x: input pytree.
    l1reg: regularization strength, float or pytree with the same structure
      as ``x``.
    scaling: a scaling factor.

  Returns:
    output pytree, with same structure as ``x``.
  """
  if l1reg is None:
    l1reg = 1.0

  fun = lambda u, v: jnp.sign(u) * jax.nn.relu(jnp.abs(u) - v * scaling)

# Path: jaxopt/_src/tree_util.py
# def broadcast_pytrees(*trees):
# def tree_scalar_mul(scalar, tree_x):
# def tree_add_scalar_mul(tree_x, scalar, tree_y):
# def _vdot_safe(a, b):
# def tree_vdot(tree_x, tree_y):
# def tree_dot(tree_x, tree_y):
# def tree_sum(tree_x):
# def tree_l2_norm(tree_x, squared=False):
# def tree_zeros_like(tree_x):
# def tree_ones_like(tree_x):
# def tree_average(trees, weights):
# def tree_gram(a):
# def tree_inf_norm(tree_x):
# def tree_where(cond, a, b):
# def tree_negative(tree):
# def tree_reciproqual(tree):
# def tree_mean(tree):


# Path: jaxopt/_src/prox.py
from typing import Any
from typing import Optional
from typing import Tuple
from jaxopt._src import tree_util
import jax
import jax.numpy as jnp

  """
  del hyperparams, scaling
  return x


def prox_lasso(x: Any,
               l1reg: Optional[Any] = None,
               scaling: float = 1.0) -> Any:
  r"""Proximal operator for the l1 norm, i.e., soft-thresholding operator.

  .. math::

    \underset{y}{\text{argmin}} ~ \frac{1}{2} ||x - y||_2^2
    + \text{scaling} \cdot \text{l1reg} \cdot ||y||_1

  When ``l1reg`` is a pytree, the weights are applied coordinate-wise.

  Args:
    x: input pytree.
    l1reg: regularization strength, float or pytree with the same structure
      as ``x``.
    scaling: a scaling factor.

  Returns:
    output pytree, with same structure as ``x``.
  """
  if l1reg is None:
    l1reg = 1.0

  fun = lambda u, v: jnp.sign(u) * jax.nn.relu(jnp.abs(u) - v * scaling)

  return tree_util.tree_multimap(fun, x, l1reg)