Update univariate.py

[marcocuturi]

correct typo

[michalk8]

Seems like the failing tests are coming from inside lineax, specifically

.tox/py3.9-jax-latest/lib/python3.9/site-packages/ott/solvers/linear/sinkhorn.py:1218: in _iterations_implicit_bwd
    out = solver.implicit_diff.gradient(
.tox/py3.9-jax-latest/lib/python3.9/site-packages/ott/solvers/linear/implicit_differentiation.py:274: in gradient
    vjp_gr = self.solve(gr, prob, f, g, lse_mode)
.tox/py3.9-jax-latest/lib/python3.9/site-packages/ott/solvers/linear/implicit_differentiation.py:220: in solve
    sch = solver(schur, res, schur_t, symmetric, **solver_kwargs)
.tox/py3.9-jax-latest/lib/python3.9/site-packages/ott/solvers/linear/lineax_implicit.py:99: in solve_lineax
    return lx.linear_solve(fn_operator, b, solver).value
.tox/py3.9-jax-latest/lib/python3.9/site-packages/lineax/_solve.py:781: in linear_solve
    operator_out_structure = operator.out_structure()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = CustomTransposeLinearOperator(
  fn=_ClosureConvert(
    jaxpr=let _where = { lambda ; a:bool[22] b:f32[22] c:i32[]. l...,), {}))),
    out_static=([], PyTreeDef(None))
  ),
  input_structure_t=ShapeDtypeStruct(shape=(22,), dtype=float32)
)

    def out_structure(self):
>       return eqxi.cached_filter_eval_shape(self.fn, self.in_structure())
E       ValueError: safe_map() argument 2 is shorter than argument 1

Will merge and look into this later, thanks @marcocuturi !