feat(useless-rewrite): Useless rewrite from log1mexp(log1mexp(x)) to x

[lmmx]

Motivation for these changes

Picking up this ticket on the PyData Global 2023 OSS sprint 🏃

• Optimize log1mexp(log1mexp(x)) -> x #471

The motivating issue concerned an issue using PyMC's Censored functionality which was traced back to the need for a "useless rewrite":

• replacing a $log(1 - exp(x))$ within another $log(1 - exp(x))$
  • i.e. a $log(1 - exp(log(1 - exp(x))))$
• by just an $x$.

Implementation details

1. 🧹 refactor(nan switch): DRY out nan switch rewrite function, easier to follow important parts 🧹

• The first change is in preparation for adding a new case, which is to simplify the case handling (avoid repetition).
• This is achieved by putting an "inner function" (function within a function) that captures the x and node from the function body in its scope, meaning we don't need to pass them as parameters, so the body of each case becomes simpler
• Every case has a "nan switch", so this trick lets us avoid repeating ourselves but retaining clarity about the variables we're using.

2. ✍️ Add new useless rewrite case ✍️

• The previous operation and the node operation are both going to be log1mexp for this case
• The condition for the case is that x >= 0 (confirm?)

Checklist

• Explain motivation and implementation 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues, preferably in nice commit messages.
• The commits correspond to relevant logical changes. Note that if they don't, we will rewrite/rebase/squash the git history before merging.
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

• N/A

New features

• "Useless rewrite" to optimise $log(1 - exp(log(1 - exp(x))))$ into just $x$

Bugfixes

• Would resolve Optimize log1mexp(log1mexp(x)) -> x #471

Documentation

• ...

Maintenance

• Took an opportunity to tidy up the 'nan switch' cases into a dict, which is clearer to read than the repetitive if blocks

[lmmx]

I think to finish this I need to add a test here

pytensor/tests/tensor/rewriting/test_math.py

Lines 1945 to 1971 in c10c376

	@pytest.mark.parametrize("exp_op", [exp, expm1])
	def test_exp_log1mexp(self, exp_op):
	# exp(log1mexp(x)) -> switch(x <= 0, 1 - exp(x), nan)
	# expm1(log1mexp(x)) -> switch(x <= 0, - exp(x), nan)
	data_valid = -np.random.random((4, 3)).astype("float32")
	data_valid[0, 0] = 0 # edge case
	data_invalid = data_valid + 1
	x = fmatrix()
	f = function([x], exp_op(log1mexp(x)), mode=self.mode)
	graph = f.maker.fgraph.toposort()
	ops_graph = [
	node
	for node in graph
	if isinstance(node.op, Elemwise)
	and isinstance(
	node.op.scalar_op, (aes.Log, aes.Log1p, aes.Log1mexp, aes.Expm1)
	)
	]
	assert len(ops_graph) == 0
	if exp_op == exp:
	expected = 1 - np.exp(data_valid)
	else:
	expected = -np.exp(data_valid)
	np.testing.assert_almost_equal(f(data_valid), expected)
	assert np.all(np.isnan(f(data_invalid)))

[ricardoV94]

This should be a bit better, as it will only call the rewrite on nodes with these Ops

Suggested change

	@node_rewriter([Elemwise])
	@node_rewriter([exp, expm1, log1mexp])

[ricardoV94]

Not needed (can't select the return None below for the git suggestion)

Suggested change

if not isinstance(node.op, Elemwise):

[ricardoV94]

This is neat but I find the old form more readable tbh? Also if I ever have to debug this rewrite I would rather have the unrolled if/elses.

I am happy with the nan_switch_helper and nesting the if/else based on the outer Ops.

[ricardoV94]

I think to finish this I need to add a test here

pytensor/tests/tensor/rewriting/test_math.py

Lines 1945 to 1971 in c10c376

	@pytest.mark.parametrize("exp_op", [exp, expm1])
	def test_exp_log1mexp(self, exp_op):
	# exp(log1mexp(x)) -> switch(x <= 0, 1 - exp(x), nan)
	# expm1(log1mexp(x)) -> switch(x <= 0, - exp(x), nan)
	data_valid = -np.random.random((4, 3)).astype("float32")
	data_valid[0, 0] = 0 # edge case
	data_invalid = data_valid + 1
	x = fmatrix()
	f = function([x], exp_op(log1mexp(x)), mode=self.mode)
	graph = f.maker.fgraph.toposort()
	ops_graph = [
	node
	for node in graph
	if isinstance(node.op, Elemwise)
	and isinstance(
	node.op.scalar_op, (aes.Log, aes.Log1p, aes.Log1mexp, aes.Expm1)
	)
	]
	assert len(ops_graph) == 0
	if exp_op == exp:
	expected = 1 - np.exp(data_valid)
	else:
	expected = -np.exp(data_valid)
	np.testing.assert_almost_equal(f(data_valid), expected)
	assert np.all(np.isnan(f(data_invalid)))

Sounds about right

[ricardoV94]

The condition for the case is that x >= 0 (confirm?)

The valid condition is that x <= 0 for which the inner log1mexp is defined. x > 0 shoud yield nan, as that would lead to taking the log of a negative number.