'Dive into Deep Learning' book cover
Dive into Deep Learning, an open source, interactive book provided in a unique form factor that integrates text, mathematics and code, recently added a new chapter on attention mechanisms.

Amazon team adds key programming frameworks to Dive into Deep Learning book

With PyTorch and TensorFlow incorporated, the authors hope to gain a wider audience.

Over the past few years, a team of Amazon scientists has been developing a book that is gaining popularity with students and developers attracted to the booming field of deep learning, a subset of machine learning focused on large-scale artificial neural networks. Called Dive into Deep Learning, the book arrives in a unique form factor, integrating text, mathematics, and runnable code. Drafted entirely through Jupyter notebooks, the book is a fully open source living document, with each update triggering updates to the PDF, HTML, and notebook versions.

Its authors are Aston Zhang, an AWS senior applied scientist; Zachary Lipton, an AWS scientist and assistant professor of Operations Research and Machine Learning at Carnegie Mellon University; Mu Li, AWS principal scientist; and Alex Smola, AWS vice president and distinguished scientist.

Recently the authors added two programming frameworks to their book: PyTorch and TensorFlow. That gives the book—originally written for MXNet—even broader appeal within the open-source machine-learning community of students, developers, and scientists. The book also is incorporated into Amazon Machine Learning University courseware.

Amazon Science spoke to the authors previously about their book, and we recently reconnected with them to learn about the significance of the new frameworks they’ve added to their book.

Q. What’s the significance of adding PyTorch and TensorFlow implementations to Dive into Deep Learning?

Mu Li: The book is designed to teach people different algorithms used in machine learning. A big asset of the book is the fact we provide all the coding information. Originally, we used MXNet because it’s a major interface and easy to learn. But then we started getting a lot of requests for PyTorch and TensorFlow implementations. So, we decided to add them to the book.

Aston Zhang: Another factor is that for machine learning practitioners, it’s not enough to know just one framework. That’s because a researcher may propose a new model or algorithm and provide implementation in only one framework. So, if you don’t know that framework, you can’t work with the model. Dive into Deep Learning now provides a way to address those different implementations. It fixes a pain point for our readers.

Zachary Lipton: Like any good product, you have to pay attention to what people are doing. And the audience available for a book that's only available in one framework is somewhat limited. Already, a great team from IIT Roorkee asked us if they could translate the code portions of our book, yielding a PyTorch version, and we gave it our blessing. We knew that a massive audience of students and practitioners would be excited for the PyTorch and TensorFlow versions.

Q. How does the change make the book better?

Alex Smola: The book is basically a collection of Jupyter notebooks – you can read the book in your web browser and run every code example in real time. Because we support multiple frameworks, we can have multiple code paths within each notebook, so you can compare both the implementations, and the results that they give side by side. That’s very powerful as a teaching tool.

Mu Li: We feel that by adding PyTorch and TensorFlow to Dive into Deep Learning, we’ve made it the best textbook to learn about and execute machine learning and deep learning. It’s a textbook, but also teaches you how to implement the code. Another thing is that some people already using PyTorch want to systematically learn deep learning. Now they can run different algorithms from scratch and learn how to do that in different frameworks.

Zachary Lipton: Nobody can survive as a professional in machine learning without having the skills to work with multiple frameworks. You might learn in MXNet or TensorFlow, but then switch jobs, and need to rapidly port those skills over to a place that uses PyTorch when you’re not familiar with it. In general, it’s important for people to learn several languages.

Q. Is any one of the frameworks superior to the others?

Alex Smola: Each of them has some advantages over the other and given the state of the open-source landscape, those advantages constantly shift. They’re all competing with each other for which is the fastest, most usable, has the best data loaders and so on. At one point, people argued that philosophy was best written in German, and music best written in Italian. If you want to have the widest audience, you don’t want to limit yourself to one approach to doing things.

Aston Zhang: We’re not asking our readers to use just one framework. We provide three implementations. Readers can click on each framework, learn how it works, and decide what works best for them. If you’re a new user, you can see the subtle differences between the three and can compare their speed. Also, we separate text and code—the text is framework-neutral, but in the code book we ask people to contribute material. We’ve had people from Google, Alibaba, IIT students and others add material. For the first few chapters, Anirudh Dagar and Yuan Tang have contributed most of the PyTorch and TensorFlow adaptations.. Many others have also helped with the adaptations to these frameworks.

Zachary Lipton: The book is starting to be useful as a Rosetta stone of sorts to allow practitioners to see what the best strategy is to solve the same problem in multiple frameworks— MXNet, PyTorch, TensorFlow—without having to chase down incompatible and idiosyncratic variants on GitHub.

Q. Was it challenging to add the different frameworks to the book?

Mu Li: Yes! PyTorch and MXNet are similar, but TensorFlow is pretty different. Fortunately, TensorFlow 2.0 is very different from TensorFlow 1.0, and closer to MXNet.

Alex Smola: The proper tuning and refinement of the models took quite a while to ensure the implementations for TensorFlow on modern convolutional neural networks were just as good as the ones in PyTorch and MXNet. That’s due to the different ways in how the frameworks implement things. And we still need to back-port the content into Mandarin. This isn’t a trivial endeavor, because there currently isn’t great tooling available to synch the text with the source code.

Q. What has been the response to the additions to Dive into Deep Learning?

Mu Li: Very good. In the past three months, compared with the prior three, we’ve seen about a 40 percent increase in users.

Q. What motivates you to continue improving Dive into Deep Learning?

Alex Smola: We write books because we want to teach and share content. It’s also our way of saying “thanks” to the machine-learning community. The book is a key enabler for spreading knowledge about machine learning and its applications much more widely. We want to make it easy for people to come in, learn about machine learning, and then surprise us with their additions to the book.

Zach Lipton: I don't think anyone involved in the project thinks of it as a book that will someday be finished, in the traditional sense. Having everything online, we can update and add material much, much more quickly than if it were made from dead trees.

Aston Zhang: Every day we get feedback from users around the world. Their comments, suggestions, encouragement, and endorsement motivates me to continue improving our book.

Research areas

Related content

IN, KA, Bengaluru
As a member of the CMT team, you'll play a key role in the evolution of our Competitive Monitoring systems to solve significantly complex and interesting technical challenges in machine learning, large language models in production, and recommender systems to name a few. The team's work directly impacts customer experience at a worldwide scale. Key job responsibilities Thought leader on the team and help set team directions Research multiple problem domains, suggest various approaches to try and be as hands-on as needed while providing more junior scientists with critical mentorship Collaborate with engineers to come up with the right LLD and HLD to solve key business problems Strong emphasis on communication via writing, internal and external talks, and being able to align with multiple stakeholders A day in the life As an Applied scientist II, a typical day will involve aligning with key product, engineering and business stakeholders ; advising junior scientists on the work they are doing ; reading current research papers and staying up-to-date on AI research ; diving deep as needed to improve CMT models and addressing stakeholders from the science perspective ; writing python code
IN, KA, Bengaluru
As a member of the CMT team, you'll play a key role in the evolution of our Competitive Monitoring systems to solve significantly complex and interesting technical challenges in machine learning, large language models in production, and recommender systems to name a few. The team's work directly impacts customer experience at a worldwide scale. Key job responsibilities 1. Research the problem domain and come up with various approaches to solve the problem. 2. Be willing to experiment quickly and fail fast. 3. Collaborate with engineers to come up with the right end to end solution to the business problems. 4. Ideate on future roadmap for science in CMT 5. Be willing to roll up your sleeves and learn core topics outside applied science, for example ML engineering A day in the life A typical day might involve (a) working on ideas for improving models around product similarity or price recommendations, (b) working closely with other scientists and our ML engineers to ensure that the best models are in production, (c) writing good maintainable code that can be reused and reproduced, (d) sharing your work across CMT and beyond via technical writings and presentations
US, CA, Santa Clara
We are looking for passionate, talented, and inventive Principal Applied Scientist with a strong machine learning background to help build industry-leading Conversational AI Systems. Our mission is to provide a delightful experience to Amazon’s customers by pushing the envelope in Natural Language Understanding (NLU), Dialog Systems including Generative AI with Large Language Models (LLMs) and Applied Machine Learning (ML). As part of our team, you will work alongside internationally recognized experts to develop novel algorithms and modeling techniques to advance the state-of-the-art in human language technology. Your work will directly impact millions of our customers in the form of products and services that make use language technology. You will gain hands on experience with Amazon’s heterogeneous text, structured data sources, and large-scale computing resources to accelerate advances in language understanding. We are hiring in all areas of human language technology: NLU, Dialog Management, Conversational AI, LLMs and Generative AI. A day in the life The team uses generative AI and foundation models to reimagine the experience of all customers on AWS. We explore new technologies and find creative solutions. Curiosity and an explorative mindset can find a place here to impact the life of engineers around the world. If you are excited about this space and want to enlighten your peers with new capabilities, this is the team for you.
US, WA, Seattle
The Automated Reasoning Group in the Amazon Neuron team is looking for an Applied Scientist to work on the intersection of Artificial Intelligence and program analysis to raise the code quality bar in our state-of-the-art deep learning compiler stack. This stack is designed to optimize application models across diverse domains, including Large Language and Vision, originating from leading frameworks such as PyTorch and JAX. Your role will involve working closely with our custom-built Machine Learning accelerator, Trainium, which represents the forefront of innovation for advanced ML capabilities, and is the underpinning of Generative AI. In this role as an Applied Scientist, you'll be instrumental in designing, developing, and deploying analyzers for ML compiler stages and compiler IRs. You will architect and implement business-critical tooling, publish research, and mentor a brilliant team of experienced scientists and engineers. You will need to be technically capable, credible, and curious in your own right as a trusted AWS Neuron engineer, innovating on behalf of our customers. Your responsibilities will involve tackling crucial challenges alongside a talented engineering team, contributing to leading-edge design and research in compiler technology and deep-learning systems software. Strong experience in programming languages, compilers, program analyzers, theorem provers, and program synthesis engines will be a benefit in this role. A background in machine learning and AI accelerators is preferred but not required.
US, CA, Sunnyvale
Prime Video is a first-stop entertainment destination offering customers a vast collection of premium programming in one app available across thousands of devices. Prime members can customize their viewing experience and find their favorite movies, series, documentaries, and live sports – including Amazon MGM Studios-produced series and movies; licensed fan favorites; and programming from Prime Video subscriptions such as Apple TV+, HBO Max, Peacock, Crunchyroll and MGM+. All customers, regardless of whether they have a Prime membership or not, can rent or buy titles via the Prime Video Store, and can enjoy even more content for free with ads. Are you interested in shaping the future of entertainment? Prime Video's technology teams are creating best-in-class digital video experience. As a Prime Video team member, you’ll have end-to-end ownership of the product, user experience, design, and technology required to deliver state-of-the-art experiences for our customers. You’ll get to work on projects that are fast-paced, challenging, and varied. You’ll also be able to experiment with new possibilities, take risks, and collaborate with remarkable people. We’ll look for you to bring your diverse perspectives, ideas, and skill-sets to make Prime Video even better for our customers. With global opportunities for talented technologists, you can decide where a career Prime Video Tech takes you! Prime Video is pioneering the use of Generative AI to empower the next generation of creatives. Our mission is to make world-class media creation accessible, scalable and efficient. We are seeking an Applied Scientist to advance the state of the art in Generative AI and to deliver these innovations as production-ready systems at Amazon scale. Your work will give creators unprecedented freedom and control while driving new efficiencies. Key job responsibilities As an Applied Scientist, you will have end-to-end ownership of the product, related research and experimentation. In addition, you will be applying advanced machine learning techniques in Computer Vision, Multimedia Understanding and Generative AI. We're building the foundational technology stack, spanning diffusion and flow-matching models, 3D/4D scene and character generation, motion and camera control, and post-training alignment. Other responsibilities include: - Research and develop generative models for controllable synthesis across images, video, vector graphics, and multimedia - Innovate in advanced diffusion and flow-based methods (e.g., inverse flow matching, parameter efficient training, guided sampling, test-time adaptation) to improve efficiency, controllability, and scalability - Advance visual grounding, depth and 3D estimation, segmentation, and matting for integration into pre-visualization, compositing, VFX, and post-production pipelines - Design multimodal GenAI workflows including visual-language model tooling, structured prompt orchestration, agentic pipelines
US, WA, Seattle
As a Principal Applied Scientist at Prime Video, you will be a technical and strategic leader responsible for inventing, developing, and deploying groundbreaking AI solutions that power personalized, relevant, and delightful experiences for millions of global customers. You will help shape the vision and direction of key ML systems that support Prime Video’s mission to deliver AI-powered customer experiences. This role demands a unique blend of deep technical expertise in machine learning and recommendation systems, industry leadership, and strong collaboration skills. You will guide the development of high-impact systems end-to-end - leading innovation from foundational research through production deployment - while mentoring scientists and influencing product and engineering roadmaps. We are looking for a thought leader who brings a strong track record of delivering ML innovations at scale, along with the curiosity and drive to push boundaries. This is a rare opportunity to drive meaningful impact at one of the largest streaming services in the world. Key job responsibilities - Invent, prototype, and productionize large-scale AI solutions across Prime Video’s personalization and discovery ecosystem using deep learning, generative AI, reinforcement learning, and optimization techniques; - Provide technical leadership and influence product vision by collaborating closely with engineers, product managers, and senior stakeholders; - Design and lead high-impact A/B tests and data analyses to validate hypotheses and guide product direction; - Drive technical bar-raising across science and engineering teams through mentorship, design reviews, and collaboration; - Stay ahead of industry trends and emerging research; leverage them to evolve long-term strategy and architecture; - Publish impactful research internally and externally (e.g. top-tier conferences and journals).
IN, KA, Bengaluru
Amazon Devices is an inventive research and development company that designs and engineer high-profile devices like the Kindle family of products, Fire Tablets, Fire TV, Health Wellness, Amazon Echo & Astro products. This is an exciting opportunity to join Amazon in developing state-of-the-art techniques that bring Gen AI on edge for our consumer products. We are looking for exceptional scientists to join our Applied Science team and help develop the next generation of edge models, and optimize them while doing co-designed with custom ML HW based on a revolutionary architecture. Work hard. Have Fun. Make History. Key job responsibilities What will you do? - Quantize, prune, distill, finetune Gen AI models to optimize for edge platforms - Fundamentally understand Amazon’s underlying Neural Edge Engine to invent optimization techniques - Analyze deep learning workloads and provide guidance to map them to Amazon’s Neural Edge Engine - Use first principles of Information Theory, Scientific Computing, Deep Learning Theory, Non Equilibrium Thermodynamics - Train custom Gen AI models that beat SOTA and paves path for developing production models - Collaborate closely with compiler engineers, fellow Applied Scientists, Hardware Architects and product teams to build the best ML-centric solutions for our devices - Publish in open source and present on Amazon's behalf at key ML conferences - NeurIPS, ICLR, MLSys.
US, MA, Boston
Are you excited about applying machine learning and applied mathematics to real-world systems at massive scale? As an Applied Scientist on this newly formed team, you will collaborate closely with scientists and engineers to bring research into production across a broad portfolio of problems — from computer vision perception platforms to building-wide optimization and orchestration. You will frame ambiguous business problems as tractable scientific challenges and implement novel machine learning (ML) systems, first-principles models, embedded systems prototypes, and performance optimizations in both prototype and production environments. This is a ground-floor opportunity to shape the scientific direction of a new organization, where your contributions will directly influence how Amazon's fulfillment network operates and evolves. Key job responsibilities - Design, develop, and deploy ML and scientific solutions spanning classical machine learning, statistical modeling, computer vision, optimization, and physics-informed modeling in production environments. - Rapidly ramp on unfamiliar problem domains, frame ambiguous business problems as tractable scientific challenges, and prototype solutions end to end. - Author or co-author research findings for internal or external peer-reviewed venues, and provide peer feedback on research procedures and results across teams. - Prototype and evaluate sensing hardware and lightweight, edge-deployable models that run on commodity compute under real-world constraints. - Collaborate across multiple science and engineering teams to integrate your solutions into deployment architecture, mentoring less experienced scientists along the way. A day in the life You might start your morning reviewing experiment results from an overnight model training run, then shift into a design discussion with engineers on how to deploy a new computer vision model to edge hardware in a fulfillment center. After lunch, you could be prototyping a physics-informed optimization approach, writing up findings for a research paper, or pairing with a teammate to debug a tricky data pipeline. As part of a new and growing organization, you will have a direct hand in shaping team practices, scientific roadmaps, and the tools you use every day. About the team Our team sits within Amazon's fulfillment technology organization and applies a range of scientific disciplines — including computer vision, optimization, reinforcement learning, and statistical modeling — to improve how goods move through Amazon's global fulfillment network. We build the models and systems that drive real-time orchestration, optimizing throughput, flow, and operational performance at scale. As a newly formed organization, we are building our culture and scientific agenda from the ground up. You will join a collaborative, inclusive group of scientists and engineers who value experimentation, rigorous research, and delivering measurable impact for customers.
CN, 31, Shanghai
Team & Project Overview The NBS Data Central team powers analytics, data science, and AI capabilities for Worldwide Global Selling (WWGS). We build scalable data products, and insight-generation systems that drive seller growth across 10+ marketplaces. Seller Intelligence is a P0 foundation theme at the Global Selling level, formed by merging "One Tagging" and "Good Contact" workstreams. It provides seller identity, segmentation, and contact-reach infrastructure that underpins all downstream seller-facing AI workflows — including intelligent outreach, personalized recommendations, and automated engagement. Scope of Impact Own the science pillar for Seller Intelligence within a cross-functional POD (PM + DE + DS + SDE) Directly impact seller engagement metrics across CN, IN, LATAM, and East-Asia expansion regions Models and data products consumed by 5+ downstream teams (ESM, NSR, MKT, NBS AI Ops, ROC) Influence $100M+ annual seller GMS through improved segmentation and contact optimization Key job responsibilities Design and deliver seller segmentation and propensity models at scale — incorporating GMS, category, growth trajectory, engagement signals, and lifecycle stage. Build contact quality scoring and lifecycle management systems (coverage optimization, dormancy detection, reactivation modeling). Define success metrics, experimentation frameworks (A/B, causal inference), and measurement methodology for seller engagement interventions. Productionize ML models and data products — partner with engineering to deploy seller scores, contact quality indices, and recommendation signals. Explore LLM/GenAI applications: automated insight generation from seller data, contact intent classification, and intelligent report synthesis. Serve as the science representative in bi-weekly NBS theme reviews; present findings and proposals to theme Bar Raisers and leadership. Collaborate with BIE team members to democratize analytical outputs via dashboards and self-serve tools. Contribute to cross-marketplace seller behavior analysis supporting Global Expansion strategy (IN, KR, VN, LATAM). Evaluate, integrate, and iterate on AI systems — assess new AI/ML tools, frameworks, and third-party models for applicability to seller intelligence use cases.
US, CA, Palo Alto
The Amazon Search team creates customer-focused search solutions and technologies. Whenever a customer visits an Amazon site worldwide and types in a query or browses through product categories, Amazon Product Search services go to work. We design, develop, and deploy high-performance distributed search systems that rank a catalog of billions of products for hundreds of millions of shoppers. The Search Relevance team owns the ranking models that decide the order of results on every Amazon search page. In this role, you will design and post-train deep ranking models, including LLM-based rankers and multi-tower deep learning models, that jointly optimize purchase, relevance, and personalization. You will invent modeling and training techniques that push the Pareto frontier across multiple objectives, and take your work end to end from novel research prototype through offline evaluation to production online experimentation. Personalization is a first-class objective on this team. You will build models that reason over each customer's history, durable preferences, and query intent to decide which results best fit that specific customer, rather than optimizing a single population-level ranking. We treat search as an active research frontier and invest heavily in staying at the leading edge of ML. Beyond today's ranking stack, our current explorations include LLM agents that reason and plan across multi-step workflows, tool-augmented foundation models, and new paradigms that combine retrieval, reasoning, and personalization. You will help chart where search goes next, and see your ideas ship to real customers within weeks, not quarters. You will work in a dynamic, entrepreneurial team while leveraging the resources of Amazon.com, one of the world's leading technology companies. Please visit https://www.amazon.science for more information. Key job responsibilities Your responsibilities include but are not limited to: - Design, train, and deploy state-of-the-art ranking models that decide how results are ordered on Amazon search, spanning LLM-based rankers and multi-tower deep learning architectures that jointly model engagement, relevance, and personalization. - Post-train LLMs and ranking models with supervised fine-tuning, reinforcement learning (e.g. GRPO, DPO, RLHF), knowledge distillation, and listwise ranking losses (e.g. LambdaLoss, ListNet, ListMLE). - Compose multiple objectives (engagement, relevance, personalization) into a single ranking through principled multi-objective optimization at inference. - Design large-scale label pipelines, including LLM-as-teacher supervision, that turn customer signals and expert judgment into training and reward signals. - Optimize inference for production ranking models through quantization, quantization-aware training, teacher-student distillation, and serving-stack tuning. - Evaluate proposed solutions through offline benchmarks and online A/B tests, and drive the analysis that decides whether a change ships. - Publish and present your work at internal and external scientific venues in ML, NLP, and IR.