remove backends folder

Author: jjcmoon

src/klay/__init__.py

@@ -15,9 +15,9 @@ def to_torch_module(self: Circuit, semiring: str = "log", probabilistic: bool =
         This means the inputs to a sum node are multiplied by a probability, and
         we can interpret sum nodes as latent Categorical variables.
     """
-    from .backends import torch
+    from .torch import KnowledgeModule
     indices = self._get_indices()
-    return torch.KnowledgeModule(*indices, semiring=semiring, probabilistic=probabilistic)
+    return KnowledgeModule(*indices, semiring=semiring, probabilistic=probabilistic)
 
 
 def to_jax_function(self: Circuit, semiring: str = "log"):
@@ -27,9 +27,9 @@ def to_jax_function(self: Circuit, semiring: str = "log"):
     :param semiring:
         The semiring in which the circuit should be evaluated. Supported options are ("log", "real", "mpe", "godel").
     """
-    from .backends import jax
+    from .jax import create_knowledge_layer
     indices = self._get_indices()
-    return jax.create_knowledge_layer(*indices, semiring=semiring)
+    return create_knowledge_layer(*indices, semiring=semiring)
 
 
 def add_sdd(self: Circuit, sdd: "SddNode", true_lits: Sequence[int] = (), false_lits: Sequence[int] = ()):

src/klay/backends/__init__.py

@@ -2,7 +2,7 @@
 import jax
 import jax.numpy as jnp
 
-from klay.backends.jax.semiring import get_semiring, encode_input
+from klay.jax.semiring import get_semiring, encode_input
 
 
 def create_knowledge_layer(pointers, ix_outs, semiring):

src/klay/backends/jax/__init__.py renamed to src/klay/jax/__init__.py

@@ -1,6 +1,6 @@
-from klay.backends.jax.semiring.godel import max_layer, min_layer
-from klay.backends.jax.semiring.log import log_sum_layer, encode_input_log
-from klay.backends.jax.semiring.real import sum_layer, prod_layer, encode_input_real
+from klay.jax.semiring.godel import max_layer, min_layer
+from klay.jax.semiring.log import log_sum_layer, encode_input_log
+from klay.jax.semiring.real import sum_layer, prod_layer, encode_input_real
 
 
 def get_semiring(name: str):

src/klay/backends/jax/semiring/__init__.py renamed to src/klay/jax/semiring/__init__.py

@@ -12,7 +12,7 @@
     pass
 
 
-from klay.backends.torch import log1mexp
+from klay.torch import log1mexp
 
 try:
     from pysdd.iterator import SddIterator