hand

venv/lib/python3.12/site-packages/jax/interpreters/partial_eval.py

@@ -0,0 +1,30 @@
+# Copyright 2018 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from jax._src.interpreters.partial_eval import (
+  DynamicJaxprTracer as DynamicJaxprTracer,
+  JaxprTracer as JaxprTracer,
+  PartialVal as PartialVal,
+  Val as Val,
+  custom_partial_eval_rules as custom_partial_eval_rules,
+  dce_jaxpr as dce_jaxpr,
+  dce_jaxpr_call_rule as dce_jaxpr_call_rule,
+  dce_jaxpr_closed_call_rule as dce_jaxpr_closed_call_rule,
+  dce_jaxpr_consts as dce_jaxpr_consts,
+  dce_rules as dce_rules,
+  partial_eval_jaxpr_custom_rules as partial_eval_jaxpr_custom_rules,
+  trace_to_jaxpr_dynamic as trace_to_jaxpr_dynamic,
+  trace_to_jaxpr_nounits as trace_to_jaxpr_nounits,
+)