pymc3 Heaviside step function

Question

In pymc3, how can the Heaviside step function be applied as a deterministic transform?

Any step function will work, but I'm asking about the Heaviside step function specifically to disambiguate it from what the documentation refers to as a "step function", which is talking about sampling steps, which are irrelevant to what I'm looking for and makes finding relevant documentation difficult to locate, if it exists. Also I wasn't able to find somewhere that all supported transforms are documented.

I want 1 if x >= 0 else 0 , but writing it literally like that won't work because presumably x >= 0 evaluates to an object, and so the if expression will immediately evaluate to 1, and then h will always be 1.

import pymc3 as pm

with pm.Model() as model:
    x = pm.Normal('x', mu=0, sigma=1)
    h = pm.Deterministic('h', 1 if x >= 0 else 0)  # ???

Answer 1

Try using theano.tensor.switch . That is,

import pymc3 as pm
import theano.tensor as tt

with pm.Model() as model:
    x = pm.Normal('x', mu=0, sigma=1)
    h = pm.Deterministic('h', tt.switch(x < 0, 0, 1))