"patch jax optional import bug; add dynamic versioning"

Author: apoorvalal

.gitignore

@@ -164,3 +164,4 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 .DS_Store
+.claude/settings.local.json

meson.build

@@ -1,6 +1,5 @@
 project('pyensmallen',
         ['cpp'],
-        version: '0.2.9',
         default_options: ['cpp_std=c++14'])
 
 py = import('python').find_installation(pure: false)

pyensmallen/__init__.py

@@ -17,9 +17,8 @@
 
 from ._pyensmallen import *
 from .losses import linear_obj, logistic_obj, poisson_obj
-from .gmm import EnsmallenEstimator
-
 
+# Core optimizers and loss functions (always available)
 __all__ = [
     "L_BFGS",
     "FrankWolfe",
@@ -32,5 +31,12 @@
     "linear_obj",
     "logistic_obj",
     "poisson_obj",
-    "EnsmallenEstimator",
 ]
+
+# Conditionally import GMM functionality (requires JAX)
+try:
+    from .gmm import EnsmallenEstimator
+    __all__.append("EnsmallenEstimator")
+except ImportError:
+    # JAX not available - install with: pip install pyensmallen[gmm]
+    pass

pyensmallen/gmm.py

@@ -9,19 +9,10 @@
 
 from . import _pyensmallen as pe
 
-try:
-    import jax
-    import jax.numpy as jnp
+import jax
+import jax.numpy as jnp
 
-    jax.config.update("jax_enable_x64", True)
-except ImportError:
-    warnings.warn(
-        "JAX is not installed. JAX autodifferentiation will not be available. "
-        "Install JAX to enable autodiff support for GMM estimation.",
-        ImportWarning,
-    )
-    jax = None
-    jnp = None
+jax.config.update("jax_enable_x64", True)
 
 
 class EnsmallenEstimator:
@@ -41,6 +32,7 @@ def __init__(
             moment_cond: Function that computes moment conditions. Should be JAX-compatible.
             weighting_matrix: Either "optimal" for two-step GMM or a custom weighting matrix
         """
+        
         self.moment_cond = moment_cond
         self.weighting_matrix = weighting_matrix
 

pyproject.toml

@@ -5,6 +5,7 @@ requires = [
     "pybind11>=2.4",
     "meson-python",
     "numpy",
+    "setuptools-scm",
 ]
 build-backend = "mesonpy"
 
@@ -13,7 +14,7 @@ skip = ["cp36-*", "*-win32", "*-manylinux_i686"]
 
 [project]
 name = 'pyensmallen'
-version = '0.2.9'
+dynamic = ["version"]
 description = 'Python bindings for the Ensmallen library.'
 readme = 'README.md'
 requires-python = '>=3.10'
@@ -22,3 +23,9 @@ authors = [
     { name = 'Apoorva Lal', email = 'lal.apoorva@gmail.com' },
     { name = 'Matthew Wardrop', email = 'mpwardrop@gmail.com' },
 ]
+
+[project.optional-dependencies]
+gmm = ["jax", "jaxlib"]
+
+[tool.setuptools_scm]
+write_to = "pyensmallen/_version.py"

setup.py

@@ -6,7 +6,6 @@
 from setuptools.command.build_ext import build_ext
 from setuptools.dist import Distribution
 
-__version__ = "0.2.9"
 
 
 class BinaryDistribution(Distribution):
@@ -99,14 +98,14 @@ def build_extensions(self):
 
 setup(
     name="pyensmallen",
-    version=__version__,
     author="Apoorva Lal",
     author_email="lal.apoorva@gmail.com",
     description="Python bindings for the ensmallen optimization library",
     long_description="",
     ext_modules=ext_modules,
     install_requires=["pybind11>=2.4"],
-    setup_requires=["pybind11>=2.4"],
+    setup_requires=["pybind11>=2.4", "setuptools-scm"],
+    use_scm_version=True,
     cmdclass={"build_ext": BuildExt},
     packages=find_packages(),
     include_package_data=True,