Moved tree_inference

[patrick-kidger]

Renamed eqx.tree_inference -> eqx.nn.inference_mode, as it's really an eqx.nn thing, not an eqx thing.

[dlwh]

while you're in here, could I avail upon you to make this work with partitioned trees?

Suggested change

	return tree_at(_inferences, pytree, replace_fn=lambda _: value)
	return tree_at(_inferences, pytree, replace_fn=lambda _: value, is_leaf: lambda x: x is None)

The issue is something like:

lora_model, base_model = eqx.partition(model, is_lora_param)  # model has dropouts
lora_model = eqx.tree_inference(lora_model, False)

this currently blows up (assuming dropouts aren't lora_params). Imho tree_inference/inference_mode should be pretty "safe" at all times?

[patrick-kidger]

Hmm. That seems dangerous, I think. If you then do eqx.combine(base_model, lora_model) -- note the order of arguments -- then the old inference value of base_model will take priority over the new inference value recorded in lora_model. You won't actually end up with an inference model.

[dlwh]

fair enough!