1 to 25 of 65 Permanent JAX Jobs

pipelines and ML models using modern libraries across multiple domains: Data science stack: NumPy, pandas/polars, scikit-learn, XGBoost, LightGBM Deep learning: PyTorch, JAX Statistical programming: NumPyro, PyMC Data skills: Proficient in SQL, with the ability to write efficient, maintainable queries and manage data pipelines for analytics and modeling ...

pipelines and ML models using modern libraries across multiple domains: Data science stack: NumPy, pandas/polars, scikit-learn, XGBoost, LightGBM Deep learning: PyTorch, JAX Statistical programming: NumPyro, PyMC Data skills: Proficient in SQL, with the ability to write efficient, maintainable queries and manage data pipelines for analytics and modeling ...

pipelines and ML models using modern libraries across multiple domains: Data science stack: NumPy, pandas/polars, scikit-learn, XGBoost, LightGBM Deep learning: PyTorch, JAX Statistical programming: NumPyro, PyMC Data skills: Proficient in SQL, with the ability to write efficient, maintainable queries and manage data pipelines for analytics and modelling ...

pipelines and ML models using modern libraries across multiple domains: Data science stack: NumPy, pandas/polars, scikit-learn, XGBoost, LightGBM Deep learning: PyTorch, JAX Statistical programming: NumPyro, PyMC Data skills: Proficient in SQL, with the ability to write efficient, maintainable queries and manage data pipelines for analytics and modelling ...

forests, and ensemble methods Experience with generative models (GANs, VAEs, Diffusion models) Programming & Tools : Expert-level Python programming Proficiency with ML frameworks (PyTorch, TensorFlow, JAX) Experience with scikit-learn, XGBoost, LightGBM Strong software engineering practices and clean code principles Data & Computing : Experience with big data technologies (Spark, Dask ...

clearly Nice-to-have Experience in optimization, reinforcement learning, and/or large language models (LLMs) Familiarity with deep learning frameworks (e.g., TensorFlow, PyTorch, JAX) Previous research publications or submitted papers Why Join Us Gain hands-on experience in state-of-the-art AI research with real-world applications Collaborate ...

modelling, RLHF). Hands-on experience training/fine-tuning generative models (LLMs, diffusion, transformers, GANs). Strong knowledge of deep learning frameworks (PyTorch, JAX, TensorFlow). Proficiency in Python and standard ML libraries. Solid foundations in probability, optimisation, and statistics. Experience working with large-scale distributed training on GPUs ...

modelling, RLHF). Hands-on experience training/fine-tuning generative models (LLMs, diffusion, transformers, GANs). Strong knowledge of deep learning frameworks (PyTorch, JAX, TensorFlow). Proficiency in Python and standard ML libraries. Solid foundations in probability, optimisation, and statistics. Experience working with large-scale distributed training on GPUs ...

e.g., PyTorch, TensorFlow) Solid grasp of ML techniques; independent and team-oriented mindset Preferred: Experience in LLMs, reinforcement learning, or chip design Familiarity with JAX and optimization frameworks Why Join Us: Work on impactful AI research with real-world applications Collaborate with top researchers in a global team Flexible hybrid ...

Strong background in deep learning and transformer architectures Hands-on experience training or fine-tuning large models (LLMs or vision models) Proficiency with PyTorch, JAX, or TensorFlow Experience with distributed training frameworks (DeepSpeed, FSDP, Megatron, ZeRO, Ray) Strong software engineering skills - writing robust, production-grade systems Experience with GPU optimization ...

expertise in Pytorch. PREFERRED QUALIFICATIONS - Master's degree in computer science or equivalent - Experience in computer architecture - Previous software engineering expertise with Pytorch/Jax/Tensorflow, Distributed libraries and Frameworks, End-to-end Model Training. Amazon is an equal opportunity employer and does not discriminate on the basis ...

software and research teams. Requirements Strong background in machine learning algorithm development and optimisation. Proficiency in Python and modern ML frameworks such as PyTorch, JAX or TensorFlow. Experience with data modelling, preprocessing and performance benchmarking. Understanding of ML evaluation metrics and model analysis methods. Excellent communication skills and collaborative mindset. ...

software and research teams. Requirements Strong background in machine learning algorithm development and optimisation. Proficiency in Python and modern ML frameworks such as PyTorch, JAX or TensorFlow. Experience with data modelling, preprocessing and performance benchmarking. Understanding of ML evaluation metrics and model analysis methods. Excellent communication skills and collaborative mindset. ...

transformers, diffusion, retrieval-augmented systems, reinforcement learning, etc.). Strong coding skills in Python and experience with at least one major ML framework (PyTorch, JAX, TensorFlow). Ability to bridge high-level research with practical implementation. Curious, humble, and excited to work on hard technical problems with real enterprise impact. ...

models and generative AI Proven experience with text-to-image generation, multimodal AI Large-scale model training experience (distributed systems, GPUs) Proficiency in PyTorch, JAX, or similar frameworks Enterprise AI deployment experience Safety and alignment research background Multilingual/multicultural AI expertise Highly Valued: A Background at top-tier ...

models and generative AI Proven experience with text-to-image generation, multimodal AI Large-scale model training experience (distributed systems, GPUs) Proficiency in PyTorch, JAX, or similar frameworks Enterprise AI deployment experience Safety and alignment research background Multilingual/multicultural AI expertise Highly Valued: A Background at top-tier ...

models and generative AI Proven experience with text-to-image generation, multimodal AI Large-scale model training experience (distributed systems, GPUs) Proficiency in PyTorch, JAX, or similar frameworks Enterprise AI deployment experience Safety and alignment research background Multilingual/multicultural AI expertise Highly Valued: A Background at top-tier ...

models and generative AI Proven experience with text-to-image generation, multimodal AI Large-scale model training experience (distributed systems, GPUs) Proficiency in PyTorch, JAX, or similar frameworks Enterprise AI deployment experience Safety and alignment research background Multilingual/multicultural AI expertise Highly Valued: A Background at top-tier ...

models and generative AI Proven experience with text-to-image generation, multimodal AI Large-scale model training experience (distributed systems, GPUs) Proficiency in PyTorch, JAX, or similar frameworks Enterprise AI deployment experience Safety and alignment research background Multilingual/multicultural AI expertise Highly Valued: A Background at top-tier ...

models and generative AI Proven experience with text-to-image generation, multimodal AI Large-scale model training experience (distributed systems, GPUs) Proficiency in PyTorch, JAX, or similar frameworks Enterprise AI deployment experience Safety and alignment research background Multilingual/multicultural AI expertise Highly Valued: A Background at top-tier ...

models and generative AI Proven experience with text-to-image generation, multimodal AI Large-scale model training experience (distributed systems, GPUs) Proficiency in PyTorch, JAX, or similar frameworks Enterprise AI deployment experience Safety and alignment research background Multilingual/multicultural AI expertise Highly Valued: A Background at top-tier ...

models and generative AI Proven experience with text-to-image generation, multimodal AI Large-scale model training experience (distributed systems, GPUs) Proficiency in PyTorch, JAX, or similar frameworks Enterprise AI deployment experience Safety and alignment research background Multilingual/multicultural AI expertise Highly Valued: A Background at top-tier ...

models and generative AI Proven experience with text-to-image generation, multimodal AI Large-scale model training experience (distributed systems, GPUs) Proficiency in PyTorch, JAX, or similar frameworks Enterprise AI deployment experience Safety and alignment research background Multilingual/multicultural AI expertise Highly Valued: A Background at top-tier ...

models and generative AI Proven experience with text-to-image generation, multimodal AI Large-scale model training experience (distributed systems, GPUs) Proficiency in PyTorch, JAX, or similar frameworks Enterprise AI deployment experience Safety and alignment research background Multilingual/multicultural AI expertise Highly Valued: A Background at top-tier ...

models and generative AI Proven experience with text-to-image generation, multimodal AI Large-scale model training experience (distributed systems, GPUs) Proficiency in PyTorch, JAX, or similar frameworks Enterprise AI deployment experience Safety and alignment research background Multilingual/multicultural AI expertise Highly Valued: A Background at top-tier ...