1 to 25 of 50 JAX Jobs

MAS). Advanced Python programming skills, with a strong emphasis on writing efficient, scalable, and maintainable code. Proven experience with TensorFlow/PyTorch/Jax, Scikit-learn, and MLOps workflows for training, deployment, and monitoring of ML models. Experience working with Polars and/or Pandas for high-performance data More ❯

field - Practical experience in solving complex problems in an applied environment - Hands on experience building models with deep learning frameworks like PyTorch, Tensorflow, or JAX - Experience building applications leveraging GenAI Amazon is an equal opportunity employer and does not discriminate on the basis of protected veteran status, disability, or other More ❯

or a related field, or equivalent professional experience Strong programming skills in Python, with experience in machine learning frameworks such as PyTorch, TensorFlow or JAX Hands-on experience fine-tuning LLMs and optimising hyperparameters for improved performance Skilled in building LLM-powered applications for real-time decision-making and automation More ❯

different programming languages (e.g., Python, R, C/C++), including experience with reinforcement learning libraries in Python Experience using deep learning packages (TensorFlow, PyTorch, JAX) Experience cleaning data, prototyping, and fine-tuning state of the art deep learning models Experience with statistics, machine learning, and reinforcement learning Experience with software More ❯

advanced machine learning techniques, with a focus on LLMs/genAI, optimisation, or sequential learning Strong background in Python, deep learning frameworks (TensorFlow, PyTorch, Jax), and ML model development PhD in Computer Science, Machine Learning, Applied Mathematics, or related field, or equivalent experience A history of successful scientific publications and More ❯

advanced machine learning techniques, with a focus on LLMs/genAI, optimisation, or sequential learning Strong background in Python, deep learning frameworks (TensorFlow, PyTorch, Jax), and ML model development PhD in Computer Science, Machine Learning, Applied Mathematics, or related field, or equivalent experience A history of successful scientific publications and More ❯

record in top-tier conferences (e.g., NeurIPS, ICML, ICLR, ACL, EMNLP, NAACL, ICLR). Familiarity with one or more deep learning frameworks (e.g. pytorch, jax, tensorflow, ) Experience in ML Research beyond completing a PhD (e.g. supervision, industry experience, leading academic initiatives, ). Excellent communication skills to report and present research More ❯

Familiarity in statistical methods for Machine Learning (e.g. Bayesian methods, HMMs, graphical models, dimension reduction, clustering, classification, regression techniques, etc.) Familiarity with PyTorch, TensorFlow, JAX or similar frameworks Strong coding abilities in Python Preferred Qualifications Experience in the Architecture, Engineering, and/or Construction domains, including expertise with industry-specific More ❯

regression techniques, etc.) Familiarity with Deep Learning techniques (e.g. Network architectures, regularization techniques, learning techniques, loss-functions, optimization strategies, etc.) Familiarity with PyTorch , TensorFlow , JAX or similar frameworks Strong coding abilities in Python and/or C++ Preferred Qualifications: 2D & 3D Generative AI Reinforcement Learning LLMs and Natural Language Processing More ❯

Strong background in machine learning, including experience with deep learning and/or generative models. Proficiency with modern ML frameworks (e.g., PyTorch, TensorFlow, or JAX). Strong coding skills (e.g., Python, C++) and a collaborative mindset. Familiarity with protein structure modelling, co-folding, or related structure-based methods, along with More ❯

data-driven role, with exposure to: building machine learning models and pipelines in Python, using common libraries and frameworks (e.g., NumPy, SciPy, Pandas, PyTorch, JAX), especially including deep learning applications; developing models for bespoke problem settings that involve high-dimensional data (spatiotemporal, geometric, physical); iterating on network architectures and model More ❯

with Deep Learning specifically with vision transformers and diffusion models. Expertise in C++ or Python with focus on deep learning frameworks like TensorFlow, PyTorch, JAX, or Scikit Learn PREFERRED QUALIFICATIONS Masters or PhD with software development experience in working with SLM/LLMs and fine tuning. Publication in leading journals More ❯

with Deep Learning specifically with vision transformers and diffusion models. Expertise in C++ or Python with focus on deep learning frameworks like TensorFlow, PyTorch, JAX, or Scikit Learn PREFERRED QUALIFICATIONS Masters or PhD with software development experience in working with SLM/LLMs and fine tuning. Publication in leading journals More ❯

with Deep Learning specifically with vision transformers and diffusion models. Expertise in C++ or Python with focus on deep learning frameworks like TensorFlow, PyTorch, JAX, or Scikit Learn PREFERRED QUALIFICATIONS Masters or PhD with software development experience in working with SLM/LLMs and fine tuning. Publication in leading journals More ❯

with Deep Learning specifically with vision transformers and diffusion models. Expertise in C++ or Python with focus on deep learning frameworks like TensorFlow, PyTorch, JAX, or Scikit Learn PREFERRED QUALIFICATIONS Masters or PhD with software development experience in working with SLM/LLMs and fine tuning. Publication in leading journals More ❯

with Deep Learning specifically with vision transformers and diffusion models. Expertise in C++ or Python with focus on deep learning frameworks like TensorFlow, PyTorch, JAX, or Scikit Learn PREFERRED QUALIFICATIONS Masters or PhD with software development experience in working with SLM/LLMs and fine tuning. Publication in leading journals More ❯

with Deep Learning specifically with vision transformers and diffusion models. Expertise in C++ or Python with focus on deep learning frameworks like TensorFlow, PyTorch, JAX, or Scikit Learn PREFERRED QUALIFICATIONS Masters or PhD with software development experience in working with SLM/LLMs and fine tuning. Publication in leading journals More ❯

quality ML systems. Preferred Qualifications: Master's degree or PhD in Computer Science or a related technical field. Experience with modern ML frameworks (e.g., JAX, Pytorch or TensorFlow). Experience with launching applied Machine Learning/Natural Language Processing (ML/NLP) projects. About the Job Google's software engineers More ❯

machine learning, with 2+ years of hands-on experience in LLM fine-tuning, RAG systems, and prompt engineering. Proficiency with frameworks like PyTorch, Tensorflow, JAX, Pandas, and OpenCV. A solid understanding of transformer models and their common variants, loss functions, and pitfalls. A quick learner with a structured, organized approach More ❯

by publications in reputable scientific journals or conferences. Expertise in software development and programming in Python. Experience with deep learning frameworks such as PyTorch, JAX, and/or TensorFlow. Excellent ability to communicate complex ideas clearly, both verbally and in writing, and to work effectively within a multidisciplinary team. Strong More ❯

applications to engineering teams and C-Level executives. This requires deep familiarity across the stack - compute infrastructure (Amazon EC2, Amazon EKA), ML frameworks PyTorch, JAX, orchestration layers Kubernetes and Slurm, parallel computing (NCCL, MPI), MLOPs, through to Amazon SageMaker Hyperpod, Amazon Bedrock as well as target use cases in the More ❯

conferences. Expertise in Programming: Strong proficiency in software development and programming in Python. Experience with Deep Learning Frameworks: Proficiency in frameworks such as PyTorch, JAX, and/or TensorFlow. Excellent Communication Skills: Ability to clearly communicate complex ideas, both verbally and in writing, and to work effectively within a multidisciplinary More ❯

conferences. Expertise in Programming: Strong proficiency in software development and programming in Python. Experience with Deep Learning Frameworks: Proficiency in frameworks such as PyTorch, JAX, and/or TensorFlow. Excellent Communication Skills: Ability to clearly communicate complex ideas, both verbally and in writing, and to work effectively within a multidisciplinary More ❯

experience in academia or industry • We will consider various levels of experience in relevant research areas Key Skills: • Experience with ML frameworks (PyTorch, TensorFlow, JAX) and efficient ML (incl. quantization, pruning, sparsification, distillation, etc.) • Experience with deployment on embedded/mobile devices (such as smartphones, with mobile CPU, GPU, NPU More ❯

machine learning, including deep learning and/or generative models. A working knowledge and practical experience with modern ML frameworks (e.g., PyTorch, TensorFlow, or JAX). Key Words: Machine Learning/Drug Discovery/Computational Chemistry/Structural Biology/AI Techniques/Protein-Ligand Interaction/Molecular Modelling/ More ❯