post nanobind docs fixes and some updates (#889)

* post nanobind docs fixes and some updates

* one more doc nit

* fix for stubs and latex

setup.py

@@ -134,9 +134,18 @@ class GenerateStubs(Command):
         pass
 
     def run(self) -> None:
-        subprocess.run(
-            ["python", "-m", "nanobind.stubgen", "-m", "mlx.core", "-r", "-O", "python"]
-        )
+        out_path = "python/mlx/core"
+        stub_cmd = [
+            "python",
+            "-m",
+            "nanobind.stubgen",
+            "-m",
+            "mlx.core",
+        ]
+        subprocess.run(stub_cmd + ["-r", "-O", out_path])
+        # Run again without recursive to specify output file name
+        subprocess.run(["rm", f"{out_path}/mlx.pyi"])
+        subprocess.run(stub_cmd + ["-o", f"{out_path}/__init__.pyi"])
 
 
 # Read the content of README.md