transform option for conditional / piecewise(ish) functions ?

[matpalm]

consider a simple function with a conditional like...

def y(x, m1, b1, split, m2, b2):
  if x > split:
    return m1*x + b1
  else:
    return m2*x + b2

grad for m1, b1, m2 and b2 work as expected ( though i imagine fitting a value for split is getting more into the land of probabilistic programming )

but i'm more interested in my options for jiting and v/pmaping functions like this... is there a general approach to handling functions like this?

[hawkinsp]

Use code like jnp.where(x > split, m1*x + b1, m2*x + b2) in this case. i.e., use predication rather than a conditional.

[matpalm]

yeah ok, so just follow the np approach of using vectorised ops; where, select, etc? wasn't sure if this was something that vmap would take care of in the cases where i was just doing definitions w.r.t scalars